// some generic types that are legal in MHLO. This pass legalizes TF types into
// types that are legal in MHLO. For example, TF::Qint8Type is converted to i8.
// Rewrites here should run before TF to MHLO op legalizations are run.

#include <memory>

#include "absl/log/log.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"

include "mlir/Pass/PassBase.td"

def TPURewritePass : Pass<"tf-tpu-rewrite", "mlir::ModuleOp"> {
  let summary = "Rewrites a `tf_device.cluster_func` on TPUs into TPU runtime operations.";

  let description = [{
    This pass rewrites a `tf_device.cluster_func` operation into a sequence of `tf._TPUCompileMlir`
    and `tf.TPUExecute` operations. `tf._TPUCompileMlir` contains a MLIR module that is

 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
// Rewrites computations generated by the xla.compile() Python code into
// XlaLaunch nodes.
//
// xla.compile() does two main things:
// a) marks operators that make up an XLA computation with the attribute

			i++
		case '/':
			if cs == 0 && cp == 0 {
				a = append(a, s[:i])
				s = s[i+1:]
				i = 0
				continue
			}
		}
		i++
	}
	return append(a, s)
}

// rewrite rewrites a subname to having only printable characters and no white
// space.
// From go/src/testing/match.go
func rewrite(s string) string {
	b := []byte{}
	for _, r := range s {
		switch {
		case isSpace(r):

  PopulateAllQuantizablePatterns(ctx, patterns);

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    // There are cases where no rewrites happen even if a pattern matches,
    // causing this to result in a convergence failure. Consider this as a
    // best-effort.
    module_op.emitWarning("Failed to converge pattern at QuantizePass.");
  }
}

  // graph rewrites. Returns false if the inference_type is DT_FLOAT or
  // `weight_quantization` flag is set.
  bool RunPropagationAndRewriteQuantizationPasses() const {
    return inference_type != tensorflow::DT_FLOAT && !weight_quantization;
  }

  // TODO: b/202075505 - make implicit weight type clearer
  // Whether run the passes and graph rewrites for dynamic range quantization.

  auto new_results =
      GetValueWithToken(builder, llvm::to_vector(terminator->getOperands()),
                        token, rewritten_operands, flatten_tuple);
  terminator->setOperands(new_results);
}

// Rewrites a `mhlo.if` op to receive and forward a `mhlo.token`. As If op does
// not have any operands other than the predicate, hence we implicitly capture
// the parent token. Also we use the same implicit token for use in the If op's

    const std::map<string, int>& host_compute_core, FunctionLibraryRuntime* flr,
    FunctionLibraryDefinition* fld, std::vector<string>* shape_inference_graphs,
    bool* has_outside_compilation);

// Rewrites XLA computation in `clusters` to replace outside compilation nodes
// with XlaHostCompute, and moves those outside compilations into `g`. If shapes
// of outside compilation outputs cannot be determined now, we will store shape

  // flatbuffer file.
  if (auto custom_op = dyn_cast<mlir::TFL::CustomOp>(*op)) {
    return custom_op.getCustomCode().starts_with("Flex");
  }

  // We never see TFL::IfOps in the IR -- it is flatbuffer_export that rewrites
  // them from TF::IfOps.
  if (isa<TF::IfOp>(op)) {
    return false;
  }

  // We assume all other TF operators are flex delegated.
  return op->getDialect()->getNamespace() == "tf";
}

                                 ArrayRef<int64_t> values, Location location) {
  auto cst_attr = rewriter.getI64TensorAttr(values);
  return rewriter.create<TF::ConstOp>(location, cst_attr.getType(), cst_attr);
}

// Rewrites broadcast->reshape to a reshape->broadcast that reduces
// the rank of the input and output of the broadcast.
class SimplifyBroadcastReshape : public OpRewritePattern<BroadcastToOp> {