// per device in `underlying_devices`. Implicit copies off of the device throw
// an error.
//
// All component tensors must have the same dtype. Currently they must also have
// the same shape, although this requirement may be relaxed in the future.
//
// `device_name` must not name an existing physical or custom device (see
// the documentation for TFE_RegisterCustomDevice for more information).
//

      auto device = op->getAttrOfType<StringAttr>(kDevice);
      for (auto *operand : operands)
        propagateIfChanged(operand, operand->SetDevice(device));
    } else {
      // Propagate device through other ops. These ops might have their
      // own device annotation, but that's fine. We only care about
      // where the TPUExecute ops live.
      StringAttr device;
      for (const Device *d : results) {

};

// A set of values that identifies a resource.
struct ResourceKey {
  StringRef device;
  StringRef container;
  StringRef shared_name;
};

ResourceKey GetResourceKey(TF::VarHandleOp var_handle_op) {
  ResourceKey resource_key;

  if (auto attr = var_handle_op->getAttrOfType<StringAttr>("device")) {
    resource_key.device = attr.getValue();
  }

  resource_key.container = var_handle_op.getContainer();

    Op<TFRT_GPU_Dialect, mnemonic, traits> {
}

// TODO(b/260267885): We may add a device argument when we want to support
// GPU MIG.
def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
  let summary = "Transfer a CPU tensor to device.";

  let description = [{
    Transfer a CPU tensor to device.

    Example:
      %device_tensor = gpurt.transfer_to_device %cpu_tensor
  }];

      // This attribute is used for op colocation. Since all ops are located
      // on a single device cluster, this private attribute is no longer
      // needed.
      op->removeAttr(kClassAttr);
      if (auto attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // Preserve device attribute if the op is placed on a replicated core
        // device. Device attribute is used to infer the appropriate sharding
        // within TPUs for this op.

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

// This pass adds the device attribute to every tf.Const op based on the device
// attribute of the operations that read its result. If the result of a tf.Const
// op is read by operations placed on multiple devices, then the pass will
// replicate the tf.Const op once for each device.

#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"

    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");

  // device cluster ops.
  for (auto func : module.getOps<func::FuncOp>()) {
    for (auto user : symbol_map.getUsers(func)) {
      // Populate caller-callee map.
      if (func::FuncOp caller = user->getParentOfType<func::FuncOp>())
        caller_callee_map[caller].insert(func);
      // Initialize function worklist with functions refrerenced in device
      // cluster.

#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
#include "tsl/profiler/protobuf/profiler_options.pb.h"

namespace tensorflow {
namespace profiler {

// Plugin initialization function that a device plugin must define. Returns
// a TF_Status output specifying whether the initialization is successful.
using TFInitProfilerFn = void (*)(TF_ProfilerRegistrationParams* const,
                                  TF_Status* const);

limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_COMPILER_JIT_XLA_COMPILER_OPTIONS_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_XLA_COMPILER_OPTIONS_UTIL_H_

#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"