See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/tfrt/function/function.h"

#include "absl/log/log.h"
#include "absl/strings/match.h"
#include "mlir/IR/OperationSupport.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project

  // to proactively transfer the result to host since the consumer op (or
  // function) of the result may still be on TPU.
  // TODO(b/194081364): remove this option once we unify servo TPU serving
  // result transfer behavior.
  bool tpu_transfer_result_to_host = false;
};

// Compile MLIR generated by tf.function in TF dialect into BEF.
Status CompileTFMLIRToBEF(const TfrtFunctionCompileOptions& options,

}
node {
  name: "custom_embedding_matmul_func_call"
  op: "custom_embedding_matmul"
}
library {
  function {
    signature {
      name: "custom_relu"
    }
    attr {
      key: "_implements"
      value {
        func {
          name: "tensorflow.relu"
        }
      }
    }
  }
  function {
    signature {
      name: "custom_embedding_matmul"
    }
    attr {
      key: "_implements"

  attr {
    key: "_disable_call_shape_inference"
    value {
      b: true
    }
  }
  attr {
    key: "_tpu_replicate"
    value {
      s: "cluster"
    }
  }
}
library {
  function {
    signature {
      name: "test_func_name"
      input_arg {
        name: "a_0"
        type: DT_INT32
      }
      output_arg {
        name: "a"
        type: DT_INT32
      }
    }

    signature {
      name: "foo3"
    }
  }
  function {
    signature {
      name: "foo4"
    }
  }
  function {
    signature {
      name: "foo5"
    }
  }
  function {
    signature {
      name: "foo6"
    }
  }
  function {
    signature {
      name: "foo7"
    }
  }
  function {
    signature {
      name: "foo8"
    }
  }
  function {
    signature {

  }
}
node {
  name: "func_call"
  op: "test_func_name"
  input: "x"
  input: "x"
  attr {
    key: "_disable_call_shape_inference"
    value {
      b: true
    }
  }
}
library {
  function {
    signature {
      name: "test_func_name"
      input_arg {
        name: "a_0"
        type: DT_RESOURCE
      }
      input_arg {
        name: "a_1"
        type: DT_RESOURCE
      }

        continue;
      }

      return symbol_use.getUser()
          ->emitError("exported function cannot be internally referenced")
          .attachNote(func.getLoc())
          .append("references this exported function");
    }
  }
  return success();
}

LogicalResult VerifyExportedFunc(func::FuncOp func) {
  bool reached_bound_inputs = false;
  auto module = func->getParentOfType<ModuleOp>();

  return IsPureFunction(func_op);
}

bool OpSideEffectCollector::IsPureFunction(func::FuncOp func_op) const {
  auto it = is_pure_function_.find(func_op);
  if (it == is_pure_function_.end()) {
    bool is_pure = true;
    is_pure_function_[func_op] = is_pure;  // prevent infinite recursion
    func_op->walk([&](Operation* op) {
      if (op == func_op) {
        return WalkResult::advance();
      }

  std::vector<Type> output_types = GetValueTypes(outputs);
  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, const llvm::SetVector<Operation*>& ops,

  assert(cost_map_.count(op) > 0);
  return cost_map_.lookup(op);
}

void CostAnalysis::AnalyzeArguments(mlir::func::FuncOp func_op) {
  // Use the max size among function inputs as the default size of dynamic
  // shaped tensors in the function.
  for (auto arg : func_op.getArguments()) {
    if (!mlir::isa<mlir::TensorType>(arg.getType())) continue;
    auto type = mlir::cast<mlir::TensorType>(arg.getType());