// Creates a pass that converts TPU models for CPU by removing TPU related ops
// such as TPUPartitionedCall, TPUReplicatedOp, etc. The TF quantizer does not
// work with models specifically designed for TPU, so this pass makes the input
// TPU model compatible with the TF quantizer by rewriting the TPU ops. The
// output model of this pass is expected to be ready for the TF quantizer.

namespace tensorflow {
namespace tf2xla {
namespace v2 {

using ::tensorflow::monitoring::testing::CellReader;
using ::testing::Not;
using ::testing::TestWithParam;
using tpu::FunctionToHloArgs;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;

static constexpr char kCompilationTimeStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/phase2_compilation_time";

    });
    if (replicate)
      reshard_was_inserted |= HandleReplicateOp(while_op, replicate);
  });
  if (reshard_was_inserted)
    VLOG(1) << "tf-tpu-variable-runtime-reformatting inserted at least one "
               "TPUReshardVariables";
  else
    VLOG(1) << "tf-tpu-variable-runtime-reformatting inserted no "
               "TPUReshardVariables";
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>

                        std::vector<int64_t>(),
                        "number of tasks from available TPU devices must be "
                        "'num_tasks' in 'topology' (1), got 2"),
        std::make_tuple(2, 1, TopologyWithMeshShapeAndTasks({1, 1, 1, 1}, 2, 2),
                        std::vector<int64_t>(),
                        "number of TPU devices available per task must be "

  void runOnFunction(
      func::FuncOp func,
      const TF::ResourceAliasAnalysis::Info& resource_alias_analysis);

  StringRef getArgument() const final { return "tf-tpu-dynamic-layout-pass"; }

  StringRef getDescription() const final {
    return "Inserts TPU layout ops to determine layout at run time.";
  }
};

// Checks if the input producer op is supported in this transform. Right now, we

    return absl::OkStatus();
  }

  // TFRT-TPU is used if device type is `DEVICE_TPU` and platform_info does not
  // have `xla_device_metadata`. This is used for TFRT-TPU when
  // BuildXlaDeviceCompiler() is called in GetCompilerIr(). Currently only
  // lowering to HLO is needed there and xla::LocalClient doesn't support
  // building the executable for TFRT-TPU and hence, is set to nullptr here.

  // CHECK: %[[IDENTIFY:.*]] = "tf.Identity"(%[[SUBGRAPH_0]]#1) {device = ""} : (tensor<1024x3xf32>) -> tensor<1024x3xf32>
  // CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1

  return builder->create<mlir::TF::ConcatOp>(
      location, output_type, concat_dimension_op.getOutput(), inputs);
}

// For tile sharded inputs to TPU computation, inject split op between the
// input values and TPU computation so that tiled input values are passed in
// as inputs to TPU computations. If more than one dimension is sharded, then
// a tree of connected split ops are added before tf_device.parallel_execute op.

  // of dropping the attribute. This need not be correct for models not trained
  // on TPU.
  // Extracts the StableHLO module from tf.XlaCallModuleOp if the StableHLO
  // module is serialized in it.
  pm.addPass(mlir::odml::CreateLegalizeTFXlaCallModuleToStablehloPass());

  // Preprocesses TPU-targeting StableHLO module for support in TF Quantizer.
  pm.addPass(mlir::quant::CreateConvertTpuModelToCpuPass());

        // device. Device attribute is used to infer the appropriate sharding
        // within TPUs for this op.
        // TODO(b/183598857): Use explicit sharding ops from the front-end.
        // For example, dequeue ops generated by
        // tensorflow/python/tpu/tpu_feed.py
        if (!tensorflow::IsTPUReplicatedCore(attr.getValue())) {
          inner_op->removeAttr(kDeviceAttr);
        }