// RUN: tf-opt -split-input-file -verify-diagnostics -tf-extract-tpu-copy-with-dynamic-shape-op %s | FileCheck %s

// Test that extract TPUCopyWithDynamicShape from host launch to device launch.

module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:COMPOSITE:0", "/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}} {

  static const char* const module_str =
      R"(
module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:GPU:0"]} {
  func.func @convert_cluster_func(%arg0: tensor<i32>) -> () {
    %2 = "tf_device.parallel_execute"() ({

      %3 = "tf_device.cluster_func"(%arg0) {device = "/job:localhost/replica:0/task:0/device:GPU:0", func = @func} : (tensor<i32>) -> tensor<i32>

func Major(dev uint64) uint32 {
	return uint32((dev >> 8) & 0xff)
}

// Minor returns the minor component of a DragonFlyBSD device number.
func Minor(dev uint64) uint32 {
	return uint32(dev & 0xffff00ff)
}

// Mkdev returns a DragonFlyBSD device number generated from the given major and
// minor components.
func Mkdev(major, minor uint32) uint64 {
	return (uint64(major) << 8) | uint64(minor)

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() -> () {
    func.return
  }

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"]} {
  func.func @main() {

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"]} {
  func.func @missing_num_cores_per_replica() {
    // expected-error@+1 {{requires attribute 'num_cores_per_replica'}}

11:name=systemd:/user/1000.user/c1.session
10:hugetlb:/user/1000.user/c1.session
9:perf_event:/user/1000.user/c1.session
8:blkio:/user/1000.user/c1.session
7:freezer:/user/1000.user/c1.session
6:devices:/user/1000.user/c1.session
5:memory:/user/1000.user/c1.session
4:cpuacct:/user/1000.user/c1.session
3:cpu:/user/1000.user/c1.session

// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-non-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:GPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: "tf.XlaLaunch"()

  }

  // Returns a device-to-device stream, in round-robin fashion.
  se::Stream* GetDeviceToDeviceStream();

  Status ThenExecute(Device* device, stream_executor::Stream* stream,
                     std::function<void()> func) override;

 private:
  bool UseMultipleStreams() const { return stream_ != host_to_device_stream_; }

  // The main compute stream of the device, used to synchronize the transfer

func Major(dev uint64) uint32 {
	return uint32((dev >> 8) & 0xff)
}

// Minor returns the minor component of a FreeBSD device number.
func Minor(dev uint64) uint32 {
	return uint32(dev & 0xffff00ff)
}

// Mkdev returns a FreeBSD device number generated from the given major and
// minor components.
func Mkdev(major, minor uint32) uint64 {
	return (uint64(major) << 8) | uint64(minor)