// CHECK-SAME: <{device = "tpu0"}>
    // CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B"(%[[A_OUTPUT]]) : (tensor<?xi32>) -> tensor<?xi32>
    %3 = "tf.B"(%2) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>

    // CHECK: %[[C_OUTPUT:[0-9]*]] = "tf.C"(%[[A_OUTPUT]], %[[B_OUTPUT]]) : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
    %4 = "tf.C"(%2, %3) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>

  auto status = tensorflow::DeviceNameUtils::DeviceNameToCpuDeviceName(
      tpu0_device, cpu0_device);
  if (!status.ok())
    return op->emitError()
           << "error in converting TPU0 to CPU0. The TPU device is "
           << tpu0_device;
  return mlir::success();
}

#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/status_helper.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/stream_executor/tpu/tpu_node_context.h"
#include "xla/stream_executor/tpu/tpu_platform.h"
#include "xla/stream_executor/tpu/tpu_platform_interface.h"
#include "xla/stream_executor/tpu/tpu_stream_interface.h"

// 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"()

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`

  }];

  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

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-tpu-merge-variables-with-execute %s | FileCheck %s

// Tests that the pass merges only variable reads/writes on the same device.

// CHECK-LABEL: func @merge_same_device_variables
// CHECK-SAME: %[[ARG_0:.*]]: tensor<*x!tf_type.resource<tensor<32xf32>>>
// CHECK-SAME: %[[ARG_1:.*]]: tensor<*x!tf_type.resource<tensor<64xf32>>>
// CHECK-SAME: %[[ARG_2:.*]]: tensor<*x!tf_type.resource<tensor<16xf32>>>

namespace tensorflow {
namespace tf2xla {
namespace v1 {

using ::tensorflow::tpu::FunctionToHloArgs;
using ::tensorflow::tpu::GuaranteedConsts;
using ::tensorflow::tpu::MlirToHloArgs;
using ::tensorflow::tpu::ShardingAndIndex;

auto* phase2_bridge_compilation_status =
    tensorflow::monitoring::Counter<1>::New(
        "/tensorflow/core/tf2xla/api/v1/"
        "phase2_compilation_status", /*metric_name*/

    tf_executor.fetch
  }
  func.return
}
// CHECK: "tf.opA"
// device = "/TPU:0"
// CHECK: _parallel_execution_ids = "r0:0"
// CHECK: "tf.opA"
// device = "/TPU:0"
// CHECK: _parallel_execution_ids = "r0:1"
// CHECK: "tf.opA"
// device = "/TPU:1"
// CHECK: _parallel_execution_ids = "r1:0"
// CHECK: "tf.opA"
// device = "/TPU:1"