// CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
// CHECK-SAME:   _xla_replica_id = 1
// CHECK:      "tf.A"
// CHECK-NOT:   _xla_replica_id
// CHECK:      tf_executor.fetch


// Tests tf._TPUDeviceOrdinalPlaceholder ops are replaced with explicit device
// ordinal constant values based on the first TPU core device id.
// CHECK-LABEL: func @device_ordinals
func.func @device_ordinals() {
  tf_executor.graph {

        %1 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 0} : () -> tensor<i64>
        %2 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 1} : () -> tensor<i64>
        tf_device.return
      }
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

    // CHECK-NEXT:     "tf_device.launch"
    // CHECK-DAG:        %[[PROGRAM_OUTPUT:[a-z_0-9]*]] = "tf._XlaCompileMlirPlaceholderProgramKey"
    // CHECK-DAG:        %[[DEVICE_ORDINAL:[a-z_0-9]+]] = "tf._TPUDeviceOrdinalPlaceholder"
    // CHECK:            %[[RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]])
    // CHECK-SAME:       key = "host_compute_channel_0_args"

func.func @not_invariant_ordinal_placeholder(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) {
  // CHECK: tf_device.replicate
  // CHECK:   tf._TPUDeviceOrdinalPlaceholder
  %0:2 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*xf32>) {n = 2: i32} {
    %1 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 0} : () -> tensor<i64>
    tf_device.return %1 : tensor<i64>
  }
  func.return
}

    return region && region->isProperAncestor(replicate_region);
  };

  for (Value operand : op->getOperands())
    if (!ancestor_of_replicate(operand.getParentRegion())) return false;

  // _TPUDeviceOrdinalPlaceholder implicitly depends on the replica.
  if (llvm::isa<TF::_TPUDeviceOrdinalPlaceholderOp>(op)) return false;

  bool has_replicate_operands = false;

    // expected-remark@above {{ID: 4}}
  }
  func.return
  // expected-remark@above {{ID: 6}}
  // expected-remark@above {{Sinks: {}}}
}

// -----

// Tests that `_TPUDeviceOrdinalPlaceholder` is side-effect-free.
func.func @device_ordinal_placeholder_side_effect_free(
  // expected-remark@above {{ID: 7}}
  ) {
  tf_executor.graph {
    // expected-remark@above {{ID: 5}}

    TF_StrTensor:$summary_metadata
  );

  let results = (outs);

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<2>;
}

def TF__TPUDeviceOrdinalPlaceholderOp : TF_Op<"_TPUDeviceOrdinalPlaceholder", [Pure]> {
  let summary = [{
Placeholder for a device ordinal that depends on its tf_device.replicate ancestor.
  }];

  let description = [{