// Build the graph:
  // "add" = "arg0" + "arg1"
  // "ret0" = "add"
  // "ret1" = "arg1"
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_INT32, 0);
  Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_FLOAT, 1);
  Output arg2 = ops::_Arg(s.WithOpName("arg2"), DT_INT32, 2);
  Output add = ops::Add(s.WithOpName("add"), arg0, arg0);

  auto arg0 = ops::_Arg(scope.WithOpName("a_0_arg"), DT_INT32, 0);
  auto arg1 = ops::_Arg(scope.WithOpName("b_0_arg"), DT_FLOAT, 1);
  auto arg2 = ops::_Arg(scope.WithOpName("c_0_arg"), DT_INT32, 2);
  auto arg3 = ops::_Arg(scope.WithOpName("d_0_arg"), DT_FLOAT, 3);

  auto arg4 = ops::_Arg(scope.WithOpName("u_0_arg"), DT_RESOURCE, 4);
  auto arg5 = ops::_Arg(scope.WithOpName("v_0_arg"), DT_RESOURCE, 5);

		// Atomic compare and swap.
		// arg0 = pointer, arg1 = old value, arg2 = new value, arg3 = memory.
		// if *arg0 == arg1 {
		//   *arg0 = arg2
		//   return (true, memory)
		// } else {
		//   return (false, memory)
		// }
		// MOV  $0, Rout
		// LR	(Rarg0), Rtmp
		// BNE	Rtmp, Rarg1, 3(PC)
		// SC	Rarg2, (Rarg0), Rtmp
		// BNE	Rtmp, ZERO, -3(PC)
		// MOV  $1, Rout

	// can be executed on g0. Second, it is called frequently, so will
	// benefit from this fast path.
	CMPQ	RARG0, $0
	JNE	rest
	get_tls(RARG0)
	MOVQ	g(RARG0), RARG0
	MOVQ	g_m(RARG0), RARG0
	MOVQ	m_p(RARG0), RARG0
	MOVQ	p_raceprocctx(RARG0), RARG0
	MOVQ	RARG0, (RARG1)
	RET

rest:
	// Transition from C ABI to Go ABI.
	PUSH_REGS_HOST_TO_ABI0()
	// Set g = g0.
	get_tls(R12)

}

TEST(IsSupportedByReplicatedBridge, ReplicatedModule) {
  const char* const code = R"mlir(
func.func @entry_func_1(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {
  %0 = "tf.Identity"(%arg0) {_tpu_replicate = "cluster"} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-graph-as-function -o - | FileCheck %s

node {
  name: "arg"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {
  name: "ret"
  op: "_Retval"
  input: "arg"
  attr {
    key: "T"
    value {
      type: DT_INT32

// RUN: tf-mlir-translate -mlir-to-graphdef %s -o - | FileCheck %s

func.func @main(%arg0: tensor<*x!tf_type.resource<tensor<8x1xf32>>>) -> tensor<8x1xf32> {
  %0 = tf_executor.graph {
     %outputs, %control = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<8x1xf32>>>) -> tensor<8x1xf32>
     tf_executor.fetch %outputs : tensor<8x1xf32>
  }
  func.return %0 : tensor<8x1xf32>
}

//   id0 = arg0
//   if-then-branch: [u0,   arg0, arg0]
//   if-else-branch: [arg0, arg0, arg1]
//     => first result is unknown, second and third is passthrough
//   if results    : [*,    arg0, {arg0, arg1}[
//   ID #2: read (unknown)         -> succ {5, 6)
//   ID #3: read (arg0)            -> succ {5}
//   ID #4: read({arg0,arg1})      -> succ {5,6}
//   ID #5: write(arg0)

		{name: "I64Eq", asm: "I64Eq", argLength: 2, reg: gp21, typ: "Bool"},   // arg0 == arg1
		{name: "I64Ne", asm: "I64Ne", argLength: 2, reg: gp21, typ: "Bool"},   // arg0 != arg1
		{name: "I64LtS", asm: "I64LtS", argLength: 2, reg: gp21, typ: "Bool"}, // arg0 < arg1 (signed)
		{name: "I64LtU", asm: "I64LtU", argLength: 2, reg: gp21, typ: "Bool"}, // arg0 < arg1 (unsigned)

  func.return
}

// CHECK-LABEL: func @callee0
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @callee0(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK-NOT:  "tf.Identity"
  // CHECK:      "tf.PartitionedCall"([[ARG0]])
  // CHECK-SAME: f = @callee1
  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>