SP_DeviceMemoryBase* const device_dst,
                       const SP_DeviceMemoryBase* const device_src,
                       uint64_t size, TF_Status* const status) {
    TF_SetStatus(status, TF_OK, "");
    std::memcpy(device_dst->opaque, device_src->opaque, size);
  };

  StreamExecutor* executor = GetExecutor(0);

                      ::testing::EquivToProto(original_graph_def)));
#endif
  }

  void InitGraphOptions() {
    session_options_.config = config_proto_;
    graph_optimization_pass_options_.device_set = &device_set_;
    graph_optimization_pass_options_.session_options = &session_options_;
    graph_optimization_pass_options_.graph = &graph_;
    graph_optimization_pass_options_.flib_def = flib_.get();
  }

                      const SP_DeviceMemoryBase* device_src, uint64_t size,
                      TF_Status* status);

  // Enqueues a memcpy operation onto stream, with a device destination
  // location and a host memory source, with target size `size`.
  void (*memcpy_htod)(const SP_Device* device, SP_Stream stream,
                      SP_DeviceMemoryBase* device_dst, const void* host_src,

  auto device_attr =
      func.getArgAttrOfType<mlir::StringAttr>(arg_no, kFuncDeviceAttr);
  return device_attr ? device_attr.getValue() : "";
}

// Extracts and canonicalizes the device attribute.
inline StringRef GetDeviceAttr(Operation* op) {
  auto device_attr = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  return device_attr ? device_attr.getValue() : "";
}

      // If op already has a TPU device set, simply propagate its device.
      auto device_attr = op_to_update->getAttrOfType<StringAttr>(kDeviceAttr);
      const bool has_device = device_attr && !device_attr.getValue().empty();
      if (has_device && tensorflow::IsTPUDevice(device_attr.getValue())) {
        PopulateDeviceForOpResults(*op_to_update, device_attr.getValue(),
                                   value_to_device);
        continue;

  mlir::func::registerAllExtensions(registry);
  // clang-format on
}

Status MlirFunctionOptimizationPass::Run(
    const std::string& function_name, const DeviceSet& device_set,
    const ConfigProto& config_proto,
    const FunctionOptimizationPass::FunctionOptions& function_options,
    std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,

 * set spans all the devices. So, a variable on a `COMPOSITE` will have a local
 * instance on the execute op's device.
 */
bool IsResourceMergeable(Attribute& resource_attr, Attribute& device_attr) {
  return resource_attr &&
         ((resource_attr == device_attr) ||
          (resource_attr.cast<mlir::StringAttr>().getValue().find(
               "COMPOSITE") != llvm::StringRef::npos));
}

  devices.push_back(FakeDevice::Make(kDeviceName, DEVICE_GPU));
  devices.push_back(FakeDevice::Make(kHostName, DEVICE_CPU));

  std::unique_ptr<DeviceSet> device_set(new DeviceSet());
  for (auto& device : devices) {
    device_set->AddDevice(device.get());
  }

  auto graph = std::make_unique<Graph>(OpRegistry::Global());
  SessionOptions session_options;
  session_options.config.mutable_graph_options()

    device_set.Insert(device_id);
  }

  TF_ASSIGN_OR_RETURN(jit::DeviceId result,
                      PickDeviceForXla(*device_info_cache, device_set,
                                       /*allow_mixing_unknown_and_cpu=*/true));
  VLOG(2) << "For " << function_name << " PickDeviceForXla("
          << device_info_cache->DebugString(device_set) << ") -> "

      }
    }

    auto device_attr = op.getAttrOfType<StringAttr>(kDeviceAttr);
    int64_t device_id;
    if (!device_attr) {
      device_id = kAnyDevice;
    } else if (devices.find(device_attr) != devices.end()) {
      device_id = devices[device_attr];
    } else {
      device_id = 1 + devices.size();
      devices[device_attr] = device_id;
    }