// is faster.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateXlaBroadcastPass();

// Creates a pass that identifies XLASharding ops in launch op for TPU
// computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUShardingIdentificationPass();

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_MARKOPSFOROUTSIDECOMPILATIONPASS

// RUN: tf-opt -tf-rewrite-tpu-embedding-ops %s | FileCheck %s

// CHECK-LABEL: func @recv_tpu_embedding_activations
func.func @recv_tpu_embedding_activations() -> (tensor<512x256xf32>) {
  // CHECK: %[[DATA:.*]] = "tf.XlaRecvTPUEmbeddingDeduplicationData"() <{config = {{.*}}}> : () -> tensor<!tf_type.variant>
  // CHECK: %[[RESULT:.*]] = "tf.XlaRecvTPUEmbeddingActivations"(%[[DATA]]) <{config = {{.*}}}> : (tensor<!tf_type.variant>) -> tensor<512x256xf32>

//   is_supported_by_replicated_brige - If the graph targets the replicated
//   bridge. Set it to true for replicated/partitioned graphs. e.g. replicated
//   and single-core TPU graphs. Set this to false if the graph is not
//   replicated, e.g. CPU/GPU graphs. is_in_fallback_enabled_mode - Whether this
//   was called with fallback to the non-MLIR Bridge. This is just for logging

  // for training.
  if (absl::StrContains(pass_options.default_device, "CPU")) {
    pass_options.skip_fold_transpose_in_ops = true;
  }
  pass_options.enable_optimizer = options.enable_optimizer;
  // Use TFRT TPU OpKernel for training.
  pass_options.target_tpurt = false;
  pass_options.tpu_use_core_selector = options.tpu_use_core_selector;
  pass_options.tpu_use_bundled_transfer = options.tpu_use_bundled_transfer;

// and logical core (`logical_core`).
// `replica_id` is the index of the ancestor ReplicateOp in [0, num_replicas).
// `logical_core` is the index of the TPU core in [0, num_cores_per_replica).
// `device_ordinal` is the index of the TPU core relative to its host.
LogicalResult GetDeviceOrdinal(const std::optional<DictionaryAttr>& devices,
                               const unsigned replica_id,

  return str_attr &&
         (str_attr.getValue().find("COMPOSITE") != llvm::StringRef::npos);
}
}  // namespace

// Rewrites unpartitioned resource reads and writes to partitioned resource
// reads and writes. The TPU computation from the frontend is generated in such
// a way that resource operations operate on the unpartitioned resource handle
// (from a `tf.TPUReplicatedInput`). This results in resource reads and writes

// RUN: tf-opt %s -tf-tpu-identity-pruning | FileCheck %s --dump-input=always

// Tests Identity op in cluster is pruned away.

// CHECK-LABEL: func @testIdentity
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @testIdentity(%arg0: tensor<i32>) {
  // CHECK-NOT:  "tf.Identity"
  // CHECK:      "tf_device.cluster"
  // CHECK-NEXT: tf_device.return [[ARG0]]
  %0 = "tf_device.cluster"() ({
    %1 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

#define DEBUG_TYPE "tf-executor-tpu-v1-island-inlining"

namespace mlir {
namespace tf_executor {

namespace {
constexpr llvm::StringRef kNestedModule = "_tpu_v1_compat_outlined";

#define GEN_PASS_DEF_EXECUTORTPUV1ISLANDINLININGPASS

// This transformation pass converts existing compilation and replication
// attributes into unified attributes. For example, A _tpu_replicate=X
// should be replaced with _xla_compile_device_type=TPU and
// _replication_info=X attributes by the conversion. An _XlaMustCompile=true
// should be replaced with _xla_compile_device_type with the value of device
// attribute.

#include "llvm/ADT/StringRef.h"

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
  // Tests that TPU cluster with no outside compilation does not generate parallel_execute.

  // CHECK-LABEL: func @no_outside_compilation
  func.func @no_outside_compilation() -> tensor<2xi32> {
    // CHECK-NOT: "tf_device.parallel_execute"