);

  let assemblyFormat = "attr-dict";
}

def CompileAndExecuteOp : TensorflowMlrtTpu_Op<"compile_and_execute"> {
  let summary = "tpu compile and execute operation";
  let description = [{
    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.

    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// TPUReshardVariables should be inserted even when While functions' shapes are
// different than While operand shapes. Test the whole tf-tpu-bridge because
// correct insertion of TPUReshardVariables depends on multiple passes including
// TPUVariableRuntimeReformatting, ShapeInference, WhileRegion canonicalization,
// and TPUMergeVariablesWithExecute.

// RUN: tf-opt %s -split-input-file -tf-tpu-dynamic-layout-pass | FileCheck %s

// Tests that the pass can transform non-replicated execution.

// CHECK: func @non_replicated(%[[ARG0:.*]]: tensor<*x!tf_type.resource> {tf.device = "/device:CPU:0"}) -> tensor<i32>
func.func @non_replicated(%arg0: tensor<*x!tf_type.resource> {tf.device = "/device:CPU:0"}) -> tensor<i32> {
  // CHECK: %[[COMPILE:.*]]:2 = "tf_device.launch"
  // CHECK-NEXT: "tf._TPUCompileMlir"()

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-tpu-annotate-dynamic-shape-inputs %s | FileCheck %s

// Test that annotate the inputs of the cluster func to be dynamic shaped.

module attributes {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"]} {
   func.func @main(
      %arg0: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"},

  }];
}

def TFTPUCompileAndExecuteOp : TensorflowMlrt_Op<"tf_tpu_compile_and_execute", [TF_MustExecute]> {
  let summary = "tpu compile and execute operation";
  let description = [{
    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.

    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project

namespace mlir {
namespace TFTPU {

// Creates a pass that rewrites `tf_device.launch_func` on TPUs into TPU runtime
// ops.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateTPURewritePass(
    llvm::StringRef module_name = llvm::StringRef());

// Creates a pass that adds ops which perform formatting on variables at

limitations under the License.
==============================================================================*/

include "mlir/Pass/PassBase.td"

def TPURewritePass : Pass<"tf-tpu-rewrite", "mlir::ModuleOp"> {
  let summary = "Rewrites a `tf_device.cluster_func` on TPUs into TPU runtime operations.";

  let description = [{
    This pass rewrites a `tf_device.cluster_func` operation into a sequence of `tf._TPUCompileMlir`

        // 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()) &&
            !isa<tf_device::LaunchOp>(op)) {

  }];

  let constructor = "tensorflow::tf2xla::internal::CreateVerifyClusteringPass()";
}

def TPUClusterFormationPass : Pass<"tf-tpu-cluster-formation", "ModuleOp"> {
  let summary = "Forms clusters from operations assigned to the same TPU computation";

  let description = [{
    TPU computations from the frontend are composed of a `tf.TPUReplicateMetadata`
    op, a subgraph of ops (TensorFlow Dialect) each with a matching

// RUN: tf-opt %s -split-input-file -tf-tpu-variable-runtime-reformatting| FileCheck %s

// Tests that the pass can correctly transform a training loop with 2 replicas.

!tf_res_f32 = tensor<*x!tf_type.resource<tensor<f32>>>
!tf_res_md_f32 = tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>> // Multi-dim f32

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  // CHECK-LABEL: func @main