// need/usage. File a bug to request porting over additional fusions.

// TODO(b/158265178): Support GPU-specific fusions.
// TODO(b/158266710): Support CPU MKL configurations.

#define GEN_PASS_DEF_FUSEDKERNELMATCHERPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

// Optimizes TF computations by fusing subgraphs/nodes onto more efficient

 private:
  void runOnOperation() override;
};

void UnwrapXlaCallModuleOp(TF::XlaCallModuleOp call_op,
                           SymbolTable& symbol_table) {
  // Do not inline lifted quantized functions used for fusing patterns.
  // TODO - b/310539922: Remove reference to TF/TFL utils.
  if (call_op->hasAttr(kQuantTraitAttrName)) {
    return;
  }

  auto function_name = call_op

    return "quant-prepare-lifting";
  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Apply graph optimizations such as fusing and constant folding to "
           "prepare lifting.";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, arith::ArithDialect>();
  }

                              additional_attributes.end()) {}

  // Name of the output kernel implementing the contraction fusion.
  std::string output_kernel;

  // Indices of additional arguments that will be forwarded to the fused
  // operation (e.g. forward bias vector if fusing BiasAdd operation).
  SmallVector<int, 4> additional_arguments;

  // Add additional attributes to the fused node.

        mlir::TFL::CreatePostQuantizePass(emit_quant_adaptor_ops));
  }
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::TFL::CreateOptimizeOpOrderPass());
  // Add optimization pass after quantization for additional fusing
  // opportunities.

  if (!pass_config.unfold_batch_matmul) {
    // Enable an optimization pass that transforms FC to BatchMatmul only when
    // `unfold_batch_matmul=false`.

    // Safeguard check to ensure that there is at least one quantizable op.
    if (failed(candidate_ops) || candidate_ops->empty()) return failure();

    // Rewrite the floating-point ops to the quantized version, by fusing
    // preceding dequantize ops and succeding quantize ops.
    for (Operation* candidate_op : *candidate_ops) {
      // If it is requantize op, we shouldn't rewrite this op.

// RUN: tf-tfrt-opt -verify-diagnostics -split-input-file -tfrt-fuse-tpu-compile-and-execute-ops -canonicalize %s | FileCheck %s --dump-input=fail --dump-input-filter=all

module attributes {tf_saved_model.semantics} {

// Test fusing _TPUCompileMlirOp and TPUExecuteOp into TPUCompileMlirAndExecuteOp.

// CHECK-LABEL: func private @test_fuse_tpu_ops
func.func private @test_fuse_tpu_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<*xi32> {

  func.return %7: tensor<1xf32>

// Fusing-LABEL: FusingaddRelu
// Fusing:  %[[add:[0-9].*]] = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<1xf32>
// Fusing:  %[[add1:[0-9].*]] = tfl.add %arg0, %[[add]] {fused_activation_function = "RELU"} : tensor<1xf32>
// Fusing:  %[[relu:[0-9].*]] = "tfl.relu"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>

std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsPass(
    const tensorflow::quantization::QuantizationOptions& quant_options);

// Apply graph optimizations such as fusing and constant folding to prepare
// lifting.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareLiftingPass(
    tensorflow::quantization::OpSet target_opset);

    auto users = op.getResult().getUsers();
    quantizing_ops.append(users.begin(), users.end());

    bool changed = false;
    // Rewrite the floating-point ops to the quantized version, by fusing
    // preceding dequantize ops and succeding quantize ops.
    for (Operation* quantizing_op : quantizing_ops) {
      // If it is requantize op, we shouldn't rewrite this op.