#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

#define DEBUG_TYPE "tf-extract-tpu-copy-with-dynamic-shape-op"

namespace mlir {
namespace TFTPU {

namespace {

#define GEN_PASS_DEF_EXTRACTTPUCOPYWITHDYNAMICSHAPEOPPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

state. It maps the nodes of a dataflow graph across many machines in a cluster, and within a machine across multiple computational devices, including multicore CPUs, general purpose GPUs, and custom-designed ASICs known as Tensor Processing Units (TPUs). This architecture gives flexibility to the application developer, whereas in previous “parameter server” designs the management of shared state is built into the system, TensorFlow enables developers to experiment with novel optimizations and training...

// This transformation pass converts unified compilation and replication
// attributes into legacy attributes. For example,  _replication_info=X
// and _xla_compile_device_type=TPU should be replaced with _tpu_replicate=X.
// This ensures the unified attributes not get exposed outside of the MLIR
// bridge with V1 pipeline in some cases.

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"

// RUN: tf-opt %s -tf-tpu-device-propagation | FileCheck %s

// Tests function passthrough values.

// CHECK-LABEL: func @testArgToRet
// CHECK-SAME: ({{%.+}}: tensor<i64> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"})
// CHECK-SAME: -> (tensor<i64> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"})
func.func @testArgToRet(%arg0: tensor<i64> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"}) -> tensor<i64> {

// RUN: tf-opt %s --tf-colocate-tpu-copy-with-dynamic-shape --split-input-file | FileCheck %s

// CHECK-LABEL test_tpu_execute
func.func @test_tpu_execute(%arg0: tensor<!tf_type.string>) {
     %0 = builtin.unrealized_conversion_cast to tensor<2048xi32>
     %1 = builtin.unrealized_conversion_cast to tensor<2048xi32>
     %2 = builtin.unrealized_conversion_cast to tensor<i32>
     %3 = builtin.unrealized_conversion_cast to tensor<i32>

		uintptr(unsafe.Pointer(&n)),
		uintptr(0),
		uintptr(0))
	if errno != 0 {
		return nil, errno
	}
	return &out, nil
}

func doinit() {
	cpuid, err := sysctlCPUID("machdep.cpu0.cpu_id")
	if err != nil {
		setMinimalFeatures()
		return
	}
	parseARM64SystemRegisters(cpuid.aa64isar0, cpuid.aa64isar1, cpuid.aa64pfr0)

	Initialized = true

  // of dropping the attribute. This need not be correct for models not trained
  // on TPU.
  // Extracts the StableHLO module from tf.XlaCallModuleOp if the StableHLO
  // module is serialized in it.
  pm.addPass(mlir::odml::CreateLegalizeTFXlaCallModuleToStablehloPass());

  // Preprocesses TPU-targeting StableHLO module for support in TF Quantizer.
  pm.addPass(mlir::quant::CreateConvertTpuModelToCpuPass());

        // device. Device attribute is used to infer the appropriate sharding
        // within TPUs for this op.
        // TODO(b/183598857): Use explicit sharding ops from the front-end.
        // For example, dequeue ops generated by
        // tensorflow/python/tpu/tpu_feed.py
        if (!tensorflow::IsTPUReplicatedCore(attr.getValue())) {
          inner_op->removeAttr(kDeviceAttr);
        }

    func.return
  }
  // CHECK:      tf_executor.island
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 0
}

// -----

// Tests device ordinal is not set for no tpu device.
// CHECK-LABEL: func @device_ordinal_attr_not_added_no_tpu_device
module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0"]} {

    void operator=(const Metadata&) = delete;
  };

  struct Options {
    // The device name's prefix (e.g., "/task:7")
    std::string device_name_prefix;

    // The name of the  device (e.g., "TPU")
    std::string device_name;

    // The index of the device.
    int device_ordinal = -1;

    // The name of the compilation device, also referred to as jit_device_type.
    // (e.g., "XLA_CPU_JIT");