// This transformation pass converts existing compilation and replication
// attributes into unified attributes. For example, A _tpu_replicate=X
// should be replaced with _xla_compile_device_type=TPU and
// _replication_info=X attributes by the conversion. An _XlaMustCompile=true
// should be replaced with _xla_compile_device_type with the value of device
// attribute.

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

==============================================================================*/

// 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.

    "_cluster_outlined_function_name";
// Marks a node for replication. The attribute value indicates the replication
// metadata op.
inline constexpr llvm::StringRef kReplicationInfoAttr = "_replication_info";
// Marks a node for XLA-TPU compilation. The attribute value indicates the
// associated compilation cluster and replication metadata op.
inline constexpr llvm::StringRef kTpuReplicateAttr = "_tpu_replicate";

    %6 = "tf.ReadVariableOp"(%arg9) : (tensor<!tf_type.resource<tensor<i64>>>) -> tensor<i64>

    ```

    wll be replaced by `_replication_info="cluster"` and  `_xla_compile_device_type="TPU"`.

    ```mlir

  "tf.ResourceApplyFtrl"(%var, %accum, %linear, %grad, %lr, %l1, %l2, %lr_power) {_xla_compile_device_type = "TPU", _replication_info = "cluster_train_function", device = "", multiply_linear_by_lr = false, use_locking = true} : (tensor<*x!tf_type.resource>, tensor<*x!tf_type.resource>, tensor<*x!tf_type.resource>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>,...

std::unique_ptr<OperationPass<func::FuncOp>>
CreateTPUPartitionedOpConversionPass();

std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass();

// Creates a pass that cleans up `_replication_info` attribute on operations
// that are inside a cluster.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUClusterCleanupAttributesPass();

// Creates a pass that removes Identity/IdentityN ops from a cluster.

        "tf.Yield"() : () -> ()
      }, {
        "tf.Yield"() : () -> ()
      }) {is_stateless = false} : (tensor<i1>) -> ()
      tf_device.return

			cfg, _ := getReplicationConfig(GlobalContext, bucket)
			if cfg != nil && cfg.RoleArn != "" {
				if dataUsageInfo.ReplicationInfo == nil {
					dataUsageInfo.ReplicationInfo = make(map[string]BucketTargetUsageInfo)
				}
				dataUsageInfo.ReplicationInfo[cfg.RoleArn] = BucketTargetUsageInfo{
					ReplicationFailedSize:   bui.ReplicationFailedSizeV1,
					ReplicationFailedCount:  bui.ReplicationFailedCountV1,

	ReplicaSize             uint64                           `json:"objectReplicaTotalSize"`
	ReplicaCount            uint64                           `json:"objectReplicaCount"`
	ReplicationInfo         map[string]BucketTargetUsageInfo `json:"objectsReplicationInfo"`
}

// DataUsageInfo represents data usage stats of the underlying Object API
type DataUsageInfo struct {