// 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>>>

module attributes {tf.versions = {producer = 888 : i32}} {
  // CHECK-LABEL: func @single_gpu_cluster_func
  func.func @single_gpu_cluster_func(%arg0: tensor<?xi32>) -> tensor<?xi32> {
    %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

// CHECK-LABEL: func @ifregion_then
// CHECK-SAME: (%arg0: {{.+}} {tf.device = "/TPU:0"}, %arg1: {{.+}} {tf.device = "/TPU:1"}
func.func @ifregion_then(
  %arg0: !tf_res,
  %arg1: !tf_res) {
  tf_executor.graph {
    // CHECK: tf_executor.island
    %island = tf_executor.island {
      // CHECK-NEXT: "tf.Identity"
      // CHECK-SAME: {device = "/TPU:0"}
      %id0 = "tf.Identity"(%arg0) : (!tf_res)
        -> !tf_res

  %3 = "tf.ReadVariableOp"(%0) {_xla_compile_device_type = "TPU", _replication_info = "a"} : (!tf_res) -> tensor<f32>
  %4 = "tf.ReadVariableOp"(%1) {_xla_compile_device_type = "TPU", _replication_info = "b"} : (!tf_res) -> tensor<f32>
  %5 = "tf.ReadVariableOp"(%2) {_xla_compile_device_type = "TPU", _replication_info = "c"} : (!tf_res) -> tensor<f32>
  %6 = "tf.Identity"(%3) {_xla_compile_device_type = "TPU", _replication_info = "a"} : (tensor<f32>) -> tensor<f32>

    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);

// Compiles a serialized MLIR module into XLA HLO, generates all accompanying
// metadata and stores them in CompilationResult.
absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,

    %outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {_xla_compile_device_type = "TPU", value = dense<2> : tensor<i32>} : () -> tensor<i32>
    %outputs_3, %control_4 = tf_executor.island wraps "tf.AddV2"(%outputs, %outputs_0) {_xla_compile_device_type = "TPU"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    tf_executor.fetch
  }

# Enable MLIR-Based new TPU Bridge

**MLIR-Based new TPU Bridge is an experimental feature, tread lightly.**

## For TF 1.x-Based Models

In tf.ConfigProto.Experimental, there is a knob controlling whether the new TPU
Bridge is enabled or not. You can set it by using the following example code:

```
session_config = tf.ConfigProto(
  ......
  experimental=tf.ConfigProto.Experimental(
    enable_mlir_bridge=True,
  ),
  ......
)
```

       "tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", device = "device", num_replicas = 1, topology = "topology"} : () -> ()
        %index = "tf.opA"(%arg0) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i1>) -> tensor<i32>
        %input = "tf.opB"(%arg0) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i1>) -> tensor<i32>

  tf_executor.graph {
    %1 = tf_executor.island {
     "tf.SomeTpuOp"() {_xla_compile_device_type = "TPU"} : () -> ()
      "tf.SomeOtherTpuOp"() {_xla_compile_device_type = "TPU"} : () -> ()
      tf_executor.yield
    }
    %2 = tf_executor.island {
      "tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "replicate", device = "device", num_replicas = 1, topology = "topology"} : () -> ()

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*/