// CHECK: "tf_device.launch"
// CHECK: device = "CORE_0"
// CHECK: "tf.opA"


// Tests devices are not remapped if device is not in replicate devices.
// CHECK-LABEL: func @no_override_device
func.func @no_override_device() {
  tf_executor.graph {
    %0 = tf_executor.island {
      tf_device.replicate {n = 2 : i32, devices = {CORE_0 = ["/CPU:0", "/GPU:1"]}} {
        "tf_device.launch"() ({
          "tf.opA"() : () -> ()

      op->erase();
      return WalkResult::advance();
    }

    if (!devices.has_value()) return WalkResult::advance();

    // Map aliased devices to explicit devices based on replica.
    if (auto launch = dyn_cast<tf_device::LaunchOp>(op))
      if (auto device_by_replica = devices.value().get(launch.getDevice()))
        launch->setAttr(
            kDeviceAttr,

// CHECK: "tf.opA"
// device = "CORE_0"
// CHECK: _parallel_execution_ids = "r0:1"


// Tests devices are not remapped if device is not in replicate devices.
// CHECK-LABEL: func @no_override_device
func.func @no_override_device() {
  tf_executor.graph {
    %0 = tf_executor.island {
      tf_device.replicate {n = 2 : i32, devices = {CORE_0 = ["/CPU:0", "/GPU:1"]}} {
        "tf_device.launch"() ({
          "tf.opA"() : () -> ()

The region held by this operation represents a computation that is replicated
across multiple devices. The number of replications is based on the `n`
attribute. Explicit devices can be populated in the `devices` attribute, and it
must be a mapping of device alias to list of explicit or aliased device names
from the outer scope. The device name map specifies devices on which replicated
ops inside tf_device.replicate will be executed.

// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
// CHECK:      return

// RUN: tf-opt %s -split-input-file -tf-xla-rewrite | FileCheck %s


module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:GPU:0"]} {
  // CHECK-LABEL: func.func @convert_cluster_func
  func.func @convert_cluster_func(%arg0: tensor<i32>) -> tensor<i32> {
    // CHECK: "tf.XlaLaunch"(%arg0) <{function = @func, operandSegmentSizes = array<i32: 0, 1, 0>}> : (tensor<i32>) -> tensor<i32>

            lhs_device_attr.getValue() == rhs_device_attr.getValue());
  };

  // Check if tf_executor.NextIteration.Source/tf_executor.NextIteration.Sink
  // pair has matching devices or no devices.
  if (auto source = llvm::dyn_cast<tf_executor::NextIterationSourceOp>(op)) {
    return ops_have_same_device(source, source.GetSink());
  } else if (auto sink = llvm::dyn_cast<tf_executor::NextIterationSinkOp>(op)) {

	}

	return annotationPrefix + name, nil
}

// annotationValue returns an annotation value for the given devices.
func annotationValue(devices []string) (string, error) {
	value, sep := "", ""
	for _, d := range devices {
		if _, _, _, err := parseQualifiedName(d); err != nil {
			return "", err
		}
		value += sep + d
		sep = ","
	}

	return value, nil
}

}

// -----

// Tests collecting compilation and execution devices results in an error.

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
  func.func @bad_devices() {
    // expected-error@+1 {{error in fetching TPU compilation/execution devices: no TPU_SYSTEM devices found}}

  if (execute_arg_to_outer_args.empty()) return false;

  // Extract the replicated devices.
  auto devices_attr = replicate.getDevices();
  if (!devices_attr) return false;

  auto device_map = devices_attr.value();
  llvm::SmallDenseMap<llvm::StringRef, llvm::SmallVector<StringRef, 4>> devices;
  devices.reserve(device_map.size());

  for (auto it : device_map) {
    auto device_alias = it.getName().strref();