Remove CUDA specific logic from runtime. (#9598)

This PR removes CUDA specific logic from `torch_xla/csrc/runtime`
directory, as well as uses of deleted functions and environment
variables from outside. This is in line with the CUDA deprecation that
started on release 2.8.

**Key Changes:**

- Removed environment variable `ZERO_COPY_ENABLED`, which was used to
use DLPack for moving without copying tensors from PyTorch CUDA to
PyTorch/XLA XLA:CUDA device
- Removed Python API function `_get_stream_for_cuda_device`, which was
used in `dlpack.py` for DLPack related logic on CUDA capsules
- Removed `ComputationClient::GetCudaStreamForDevice()`, which was used
by the Python API above
- Removed `PjRtComputationClient::RegisterCustomCall()`, since it only
worked when `platform == "CUDA"`
- Removed `GetGpuAllocatorConfig()`
- Removed `from_xla_cuda_to_cuda()` DLPack function
- Removed CUDA branch from `InitializePjRt()`

Author: ysiraichi

test/dynamo/test_dynamo.py

@@ -157,9 +157,6 @@ def _choose_proper_device(self, initialize_on_cuda):
       self.skipTest(
           "Skip this test because it requires xr.device_type()=='CUDA' and torch.cuda.is_available()."
       )
-    os.environ.update({
-        xenv.ZERO_COPY_ENABLED: "1",
-    })
     return "cuda:0"
 
   @skipOnNeuron
@@ -205,9 +202,6 @@ def test_simple_model(self):
       "1",
   )
   def test_simple_model_automoves_tensors(self, zero_copy_enabled):
-    os.environ.update({
-        xenv.ZERO_COPY_ENABLED: zero_copy_enabled,
-    })
     x = torch.tensor(100.0, requires_grad=True, device="cuda:0")
     y = torch.tensor(200.0, requires_grad=True, device="cuda:0")
     original_device = x.device

torch_xla/_dynamo/dynamo_bridge.py

@@ -148,19 +148,10 @@ def _maybe_move_tensors_to_device(tensors: tuple,
     if dynamo_debug:
       print("Moving Tensor {} to device {}".format(tensor, target_device))
 
-    zero_copy_enabled = xu.getenv_as(xenv.ZERO_COPY_ENABLED, bool, defval=False)
-    if zero_copy_enabled and tensor.device.type == 'cuda' and target_device.type == 'xla':
-      # If the input cuda tensor requires gradient, we need to call detach. Otherwise, we'd get the error "RuntimeError: Can't export tensors that require gradient, use tensor.detach()"
-      moved_tensor = torch_xla_dlpack.from_dlpack(tensor.detach())
-    elif zero_copy_enabled and tensor.device.type == 'xla' and target_device.type == 'cuda':
-      # `torch_xla.sync()` is need to make sure the pjrt buffer is valid.
-      torch_xla.sync()
-      moved_tensor = torch_xla_dlpack.from_xla_cuda_to_cuda(tensor)
-    else:
-      # Have to move to CPU before moving it to target device.
-      cpu_device: torch.device = torch.device("cpu")
-      moved_tensor = tensor.to(cpu_device)
-      moved_tensor = moved_tensor.to(target_device)
+    # Have to move to CPU before moving it to target device.
+    cpu_device: torch.device = torch.device("cpu")
+    moved_tensor = tensor.to(cpu_device)
+    moved_tensor = moved_tensor.to(target_device)
 
     # Explicitly have to copy requires_grad attribute because it's dropped
     # with torch.to(..)

torch_xla/core/xla_env_vars.py

@@ -30,4 +30,3 @@
 RANK = 'RANK'
 WORLD_SIZE = 'WORLD_SIZE'
 LOCAL_WORLD_SIZE = 'LOCAL_WORLD_SIZE'
-ZERO_COPY_ENABLED = 'ZERO_COPY_ENABLED'

torch_xla/csrc/init_python_bindings.cpp

@@ -1771,11 +1771,6 @@ void InitXlaModuleBindings(py::module m) {
            []() {
               return runtime::GetComputationClientOrDie()->GetPlatformVersion();
            })
-      .def("_get_stream_for_cuda_device",
-           [](const int device_id) {
-            return runtime::GetComputationClientOrDie()->GetCudaStreamForDevice(
-                device_id);
-           })
       .def("_xla_num_devices",
            []() -> int64_t {
             if (UseVirtualDevice()) {

torch_xla/csrc/runtime/computation_client.h

@@ -375,8 +375,6 @@ class ComputationClient {
   virtual absl::StatusOr<xla::PjRtDevice*> LookupAddressableDevice(
       int local_device_id) const = 0;
 
-  virtual std::intptr_t GetCudaStreamForDevice(int local_device_id) const = 0;
-
   virtual size_t GetNumLocalDevices() const = 0;
 
   virtual size_t GetNumDevices() const = 0;

torch_xla/csrc/runtime/env_vars.h

@@ -10,23 +10,19 @@ namespace env {
 inline constexpr char kEnvLocalWorker[] = "LOCAL_WORKER";
 inline constexpr char kEnvTpuConfig[] = "TPU_CONFIG";
 inline constexpr char kEnvNumTpu[] = "TPU_NUM_DEVICES";
-inline constexpr char kEnvNumGpu[] = "GPU_NUM_DEVICES";
 inline constexpr char kEnvNumCpu[] = "CPU_NUM_DEVICES";
 inline constexpr char kEnvTpuvmMode[] = "TPUVM_MODE";
 inline constexpr char kEnvPjRtDevice[] = "PJRT_DEVICE";
 inline constexpr char kEnvPjRtTpuMaxInflightComputations[] =
     "PJRT_TPU_MAX_INFLIGHT_COMPUTATIONS";
 inline constexpr char kEnvPjrtAsyncCpuClient[] = "PJRT_CPU_ASYNC_CLIENT";
-inline constexpr char kEnvPjrtAsyncGpuClient[] = "PJRT_GPU_ASYNC_CLIENT";
 inline constexpr char kEnvTpuLibraryPath[] = "TPU_LIBRARY_PATH";
 inline constexpr char kEnvInferredTpuLibraryPath[] = "PTXLA_TPU_LIBRARY_PATH";
 inline constexpr char kEnvXpuLibraryPath[] = "XPU_LIBRARY_PATH";
 inline constexpr char kEnvNeuronLibraryPath[] = "NEURON_LIBRARY_PATH";
 inline constexpr char kEnvPjrtDistServiceAddr[] = "PJRT_DIST_SERVICE_ADDR";
 inline constexpr char kEnvPjRtLocalProcessCount[] = "PJRT_LOCAL_PROCESS_COUNT";
 inline constexpr char kEnvPjRtLocalRank[] = "PJRT_LOCAL_PROCESS_RANK";
-inline constexpr char kEnvPjrtAllocatorCudaAsync[] =
-    "PJRT_ALLOCATOR_CUDA_ASYNC";
 inline constexpr char kEnvPjrtAllocatorPreallocate[] =
     "PJRT_ALLOCATOR_PREALLOCATE";
 inline constexpr char kEnvPjrtAllocatorFraction[] = "PJRT_ALLOCATOR_FRACTION";

torch_xla/csrc/runtime/ifrt_computation_client.cpp

@@ -161,8 +161,6 @@ IfrtComputationClient::Create() {
 }
 
 IfrtComputationClient::~IfrtComputationClient() {
-  // In the GPU case, the PjRtClient depends on the DistributedRuntimeClient
-  // tracked in XlaCoordinator, so the PjRtClient must be destroyed first.
   client_ = nullptr;
   coordinator_ = nullptr;
 }

torch_xla/csrc/runtime/ifrt_computation_client.h

@@ -110,10 +110,6 @@ class IfrtComputationClient : public ComputationClient {
     XLA_ERROR() << __FUNCTION__ << " not implemented";
   }
 
-  std::intptr_t GetCudaStreamForDevice(int local_device_id) const override {
-    XLA_ERROR() << __FUNCTION__ << " not implemented";
-  }
-
   std::vector<std::string> GetLocalDevices() const override;
 
   std::vector<std::string> GetAllDevices() const override;

torch_xla/csrc/runtime/pjrt_computation_client.cpp

@@ -23,7 +23,6 @@
 #include "xla/hlo/builder/xla_builder.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/literal.h"
-#include "xla/pjrt/c/pjrt_c_api_gpu_extension.h"
 #include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
 #include "xla/pjrt/pjrt_api.h"
 #include "xla/pjrt/pjrt_c_api_client.h"
@@ -152,8 +151,6 @@ PjRtComputationClient::Create() {
 }
 
 PjRtComputationClient::~PjRtComputationClient() {
-  // In the GPU case, the PjRtClient depends on the DistributedRuntimeClient
-  // tracked in XlaCoordinator, so the PjRtClient must be destroyed first.
   client_ = nullptr;
   coordinator_ = nullptr;
 }
@@ -1038,45 +1035,6 @@ ComputationClient::MemoryInfo PjRtComputationClient::GetMemoryInfo(
   };
 }
 
-void PjRtComputationClient::RegisterCustomCall(const std::string& fn_name,
-                                               void* function_ptr,
-                                               const std::string& platform) {
-  if (platform != "CUDA") {
-    XLA_ERROR() << "Custom call targets can only be registered for "
-                   "PJRT CUDA runtime.";
-    return;
-  }
-
-  auto* c_api_client = dynamic_cast<xla::PjRtCApiClient*>(client_.get());
-  if (!c_api_client) {
-    XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(fn_name, function_ptr, platform);
-    return;
-  }
-  const PJRT_Api* pjrt_api = c_api_client->pjrt_c_api();
-
-  // See openxla reference:
-  // https://github.com/openxla/xla/blob/b604c8d87df842002a7a8de79a434026329fbcb2/xla/pjrt/c/pjrt_c_api_gpu_test.cc#L414
-  const PJRT_Extension_Base* next =
-      reinterpret_cast<const PJRT_Extension_Base*>(pjrt_api->extension_start);
-  while (next != nullptr &&
-         next->type !=
-             PJRT_Extension_Type::PJRT_Extension_Type_Gpu_Custom_Call) {
-    next = next->next;
-  }
-  XLA_CHECK(next) << "Custom call extension not found";
-  PJRT_Gpu_Register_Custom_Call_Args args;
-  args.struct_size = PJRT_Gpu_Register_Custom_Call_Args_STRUCT_SIZE;
-  args.function_name = fn_name.c_str();
-  args.function_name_size = fn_name.size();
-  args.api_version = 0;
-  args.handler_execute = function_ptr;
-  PJRT_Error* error =
-      reinterpret_cast<const PJRT_Gpu_Custom_Call*>(next)->custom_call(&args);
-  if (error) {
-    XLA_ERROR() << error->status;
-  }
-}
-
 void PjRtComputationClient::OnReadyCallback(
     ComputationClient::DataPtr data, const std::function<void()>& callback) {
   std::shared_ptr<xla::PjRtBuffer> buffer;

torch_xla/csrc/runtime/pjrt_computation_client.h

@@ -118,17 +118,6 @@ class PjRtComputationClient : public ComputationClient {
         xla::PjRtLocalDeviceId(local_device_id));
   }
 
-  std::intptr_t GetCudaStreamForDevice(int local_device_id) const override {
-    absl::StatusOr<xla::PjRtDevice*> pjrt_device =
-        client_->LookupAddressableDevice(
-            xla::PjRtLocalDeviceId(local_device_id));
-    XLA_CHECK(pjrt_device.ok()) << "Failed to get a PjRt device.";
-    absl::StatusOr<std::intptr_t> stream =
-        pjrt_device.value()->GetStreamForExternalReadyEvents();
-    XLA_CHECK(stream.ok()) << "Failed to get a stream.";
-    return stream.value();
-  }
-
   std::vector<std::string> GetLocalDevices() const override;
 
   std::vector<std::string> GetAllDevices() const override;
@@ -169,7 +158,9 @@ class PjRtComputationClient : public ComputationClient {
       absl::Span<xla::PjRtDevice* const> devices) const;
 
   void RegisterCustomCall(const std::string& fn_name, void* function_ptr,
-                          const std::string& platform) override;
+                          const std::string& platform) override {
+    XLA_ERROR() << __FUNCTION__ << " not implemented";
+  };
 
   void OnReadyCallback(DataPtr data,
                        const std::function<void()>& callback) override;