// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-tpu -split-input-file -verify-diagnostics %s | FileCheck %s

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_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: %0:2 = "tf_device.launch"() <{{.*}}> ({

    SessionOptions options;
    std::vector<std::unique_ptr<Device>> devices;
    Status s = device_factory->CreateDevices(
        options, "/job:worker/replica:0/task:0", &devices);
    device_ = std::move(devices[0]);

    AllocatorAttributes host_alloc_attr;
    host_alloc_attr.set_on_host(true);
    host_allocator_ = device_->GetAllocator(host_alloc_attr);

    AllocatorAttributes device_alloc_attr;

func (info *ClaimInfo) cdiDevicesAsList() []kubecontainer.CDIDevice {
	var cdiDevices []kubecontainer.CDIDevice
	for _, devices := range info.CDIDevices {
		for _, device := range devices {
			cdiDevices = append(cdiDevices, kubecontainer.CDIDevice{Name: device})
		}
	}
	return cdiDevices
}

// addPodReference adds a pod reference to the claim info.

    Builder* builder, OperationState* state, int n,
    std::optional<DictionaryAttr> devices,
    llvm::ArrayRef<std::pair<ValueRange, Type>> replicated_inputs,
    ValueRange packed_inputs, TypeRange replica_output_types) {
  DCHECK_GE(n, 2);
  state->addAttribute("n", builder->getI32IntegerAttr(n));

  if (devices.has_value()) state->addAttribute("devices", devices.value());

  Region* region = state->addRegion();

    // TODO(jpienaar): Remove this later.
    if (auto device = res->getAttrOfType<StringAttr>("device")) {
      if (!device.getValue().empty())
        result_id->setAttr("device", device);
      else
        result_id->setAttr("device", compile_device_op);
    } else if (compile_device_op) {
      result_id->setAttr("device", compile_device_op);
    }

  auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
  // Get runtime devices information from the closest parent module.
  auto module = getOperation();
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices)))
    return signalPassFailure();

  OpBuilder builder(&getContext());
  llvm::SmallVector<mlir::tf_device::ClusterOp, 4> clusters;

  }

  return resource_users;
}

void ColocateCompositeResourceOpsInReplicate(
    tf_device::ReplicateOp replicate_op, OpBuilder* builder) {
  auto devices = replicate_op.getDevices();
  if (!devices) return;
  if (!devices.value().get(tensorflow::GetDeviceAliasForLogicalCore(0))) return;

  const auto composite_resource_users =
      GetResourceOpsUsingCompositeArgsInReplicate(replicate_op);

    // `enabled_for_gpu_` below.
    bool enabled_for_all_;

    // If true, enable Device API (PjRt) for TF GPU device. This is a helper
    // flag so that individual tests can turn on PjRt for GPU specifically.
    // Once the rollout to GPU is complete, this flag can be deprecated.
    bool enabled_for_gpu_;

   private:
    // Devices for which using Device API (PjRt) is allowed in the XlaLaunch op.

	s.Send(&pluginapi.ListAndWatchResponse{Devices: m.devs})

	for {
		select {
		case <-m.stop:
			return nil
		case updated := <-m.update:
			s.Send(&pluginapi.ListAndWatchResponse{Devices: updated})
		}
	}
}

// Update allows the device plugin to send new devices through ListAndWatch
func (m *Stub) Update(devs []*pluginapi.Device) {
	m.update <- devs
}

    const string dev_name(TF_DeviceListName(devices, i, status.get()));
    CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
    if (dev_type == device_type) {
      *device_name = dev_name;
      LOG(INFO) << "Found " << device_type << " device " << *device_name;
      TF_DeleteDeviceList(devices);
      return true;
    }
  }
  TF_DeleteDeviceList(devices);
  return false;
}