Remove CUDA logic from C++ files in torch_xla/csrc directory. (#9603)

This PR removes CUDA specific code from C++ files in `torch_xla/csrc`
directory. This is in line with the CUDA deprecation that started on
release 2.8.

**Key Changes:**

- (`init_python_bindings.cpp`) Removed `_has_cuda_support` Python API
- (`dl_convertor.cpp`) Removed CUDA handling of DLPack capsules
- (`tensor_impl.cpp`) Removed special handling of `Autocast` dispatch
key for XLA:CUDA device
- Also added a check, crashing on `XLA:CUDA` device (shouldn't be
supported anymore)

Author: ysiraichi

torch_xla/__init__.py

@@ -259,8 +259,7 @@ def _init_xla_lazy_backend():
 from .experimental import plugins
 from ._internal import neuron, xpu  # Additional built-in plugins
 
-if os.getenv('XLA_REGISTER_INSTALLED_PLUGINS',
-             '0' if _XLAC._has_cuda_support() else '1') == '1':
+if os.getenv('XLA_REGISTER_INSTALLED_PLUGINS', '1') == '1':
   plugins.use_dynamic_plugins()
   plugins.register_installed_plugins()
 

torch_xla/csrc/dl_convertor.cpp

@@ -51,8 +51,6 @@ void DLPackTensorDeleter(DLManagedTensor* t) {
 DLDeviceType DLDeviceTypeForDevice(const xla::PjRtDevice& device) {
   if (device.client()->platform_id() == xla::CpuId()) {
     return DLDeviceType::kDLCPU;
-  } else if (device.client()->platform_id() == xla::CudaId()) {
-    return DLDeviceType::kDLCUDA;
   }
   XLA_ERROR() << "Device " << device.DebugString()
               << " cannot be used as a DLPack device.";
@@ -176,11 +174,6 @@ absl::StatusOr<xla::PjRtDevice*> DeviceForDLDevice(const DLDevice& context) {
                    xla::CpuId());
       return runtime::GetComputationClientOrDie()->LookupAddressableDevice(
           context.device_id);
-    case DLDeviceType::kDLCUDA:
-      XLA_CHECK_EQ(runtime::GetComputationClientOrDie()->GetPlatformID(),
-                   xla::CudaId());
-      return runtime::GetComputationClientOrDie()->LookupAddressableDevice(
-          context.device_id);
     default:
       return tsl::errors::InvalidArgument(
           "Unknown/unsupported DLPack device type %d", context.device_type);

torch_xla/csrc/init_python_bindings.cpp

@@ -3055,14 +3055,6 @@ void InitXlaModuleBindings(py::module m) {
                -> std::vector<at::Tensor> {
             return TpuCustomCall(inputs, payload, output_shapes, output_dtypes);
            })
-      .def("_has_cuda_support",
-           []() {
-#ifdef GOOGLE_CUDA
-             return true;
-#else
-            return false;
-#endif
-           })
       .def("_xla_register_custom_call_target",
            [](const std::string& fn_name, const py::capsule& function_ptr,
               const std::string& platform) {

torch_xla/csrc/random.cpp

@@ -16,16 +16,7 @@
 namespace torch_xla {
 namespace {
 
-std::string GetDefaultGitGeneratorName() {
-  XlaDeviceType hw_type =
-      static_cast<XlaDeviceType>(bridge::GetCurrentDevice().type());
-  switch (hw_type) {
-    case XlaDeviceType::CUDA:
-      return "three_fry";
-    default:
-      return "default";
-  }
-}
+std::string GetDefaultGitGeneratorName() { return "default"; }
 
 xla::BitGeneratorTy GetBitGenerator() {
   static const std::string* bit_generator =

torch_xla/csrc/tensor_impl.cpp

@@ -9,6 +9,7 @@
 #include <torch/csrc/lazy/core/tensor_util.h>
 #include <torch/csrc/lazy/core/util.h>
 
+#include "absl/log/absl_check.h"
 #include "torch_xla/csrc/aten_xla_bridge.h"
 #include "torch_xla/csrc/device.h"
 #include "torch_xla/csrc/ir_builder.h"
@@ -71,16 +72,11 @@ XLATensorImpl::XLATensorImpl(XLATensor&& tensor)
                       GetTypeMeta(tensor),
                       bridge::XlaDeviceToAtenDevice(tensor.GetDevice())),
       tensor_(c10::make_intrusive<XLATensor>(std::move(tensor))) {
-  // Update the Autocast key based off the backend device.
-  // Upstream TensorImpl cannot differentiate between XLA:TPU and XLA:GPU
-  // so we must manually update Autocast to AutocastCUDA on XLA:GPU.
-  torch::lazy::BackendDevice current_device = bridge::GetCurrentDevice();
-  auto dev_type = static_cast<XlaDeviceType>(current_device.type());
-  if (dev_type == XlaDeviceType::CUDA) {
-    auto autocast_cuda_ks = c10::DispatchKeySet(c10::DispatchKey::AutocastCUDA);
-    auto autocast_xla_ks = c10::DispatchKeySet(c10::DispatchKey::AutocastXLA);
-    key_set_ = (key_set_ - autocast_xla_ks) | autocast_cuda_ks;
-  }
+  auto dev_type = static_cast<XlaDeviceType>(bridge::GetCurrentDevice().type());
+  ABSL_CHECK(dev_type != XlaDeviceType::CUDA)
+      << "XLA:CUDA is not supported anymore. "
+         "If you are seeing this error, report a bug to the PyTorch/XLA GitHub "
+         "repository: https://github.com/pytorch/xla";
   const_cast<XLATensorImpl*>(this)->SetupSizeProperties();
   set_sizes_and_strides(sym_sizes_, c10::fromIntArrayRefSlow(
                                         sizes_and_strides_.strides_arrayref()));