#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/lib/core/status_test_util.h"

namespace mlir::quant::stablehlo {
namespace {

using ::testing::Test;

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

constexpr char kEntryFuncName[] = "main";

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

using ::tensorflow::monitoring::testing::CellReader;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;

  return absl::OkStatus();
}

// TODO(b/305211853): Unify the CPU/TPU/GPU Execution Ops and thus these two
// passes should merge together.
void RegisterTPULowerClusterToRuntimeOpsPassPipeline() {
  static mlir::PassPipelineRegistration<StandardPipelineOptions> pipeline(
      "tfrt-lower-cluster-to-runtime-ops-tpu",
      "Run all the passes involved after the clustering transformations from "

source ci/official/envs/linux_x86
TFCI_BAZEL_COMMON_ARGS="--repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config release_cpu_linux --config=tpu"
TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu
TFCI_BUILD_PIP_PACKAGE_ARGS="--repo_env=WHEEL_NAME=tensorflow_tpu"
TFCI_LIB_SUFFIX="-tpu-linux-x86_64"
TFCI_WHL_BAZEL_TEST_ENABLE=0
TFCI_WHL_IMPORT_TEST_ENABLE=0
TFCI_WHL_SIZE_LIMIT=580M

  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) {
    cluster_funcs.push_back(cluster_func);
  });

  OpBuilder builder(&getContext());
  // Add resource reads for resource writes from TPU cluster where for such
  // resources the TPU cluster does not read from.
  for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) {
    builder.setInsertionPoint(cluster_func);

    SmallVector<Value, 4> read_operands;

#include <variant>
#include <vector>

#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Fills in arg_shapes and metadata_proto with appropriate values based on the
// input mlir module.

                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and
    // H2D transfer kernels. So fold the tf.Identity op if:
    // * the identity op is placed on TPU, and
    // * the arg to the identity op is produced by a TPUExecuteOp.
    if (op_device && op_device.getValue().contains("TPU")) {
      return true;
    }

    Operation *def_op = arg.getDefiningOp();

    compatible_with = get_compatible_with_portable(),  # copybara: comment
    visibility = [
        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/runtime_fallback:internal",
        "//tensorflow/core/tfrt/saved_model:friends",

  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}

// Adds outside compilation attributes to unary ops such as Identity/Cast ops
// at the head of TPU computation that is used only by other outside compiled
// ops. Identity ops and Cast ops is commonly added to the start of TPU
// computation. Adding/expanding outside compilation attributes to these ops
// will ensure that head outside compiled ops are correctly located and moved to
// host.