Remove CUDA from PyTorch/XLA build. (#9609)

This PR removes the CUDA specific build related source code. This is in
line with the CUDA deprecation that started on release 2.8.

**Key Changes:**

- (`.bazelrc`) Removed CUDA specific environment variables
- (`WORKSPACE`) Removed CUDA dependencies and CUDA specific OpenXLA
patches
- (`BUILD` files) Removed CUDA specific build dependencies

Author: ysiraichi

.bazelrc

@@ -91,7 +91,7 @@ build:short_logs --output_filter=DONT_MATCH_ANYTHING
 #build:tpu --@xla//xla/python:enable_tpu=true
 build:tpu --define=with_tpu_support=true
 
-# Run tests serially with TPU and GPU (only 1 device is available).
+# Run tests serially with TPU (only 1 device is available).
 test:tpu --local_test_jobs=1
 
 #########################################################################
@@ -100,11 +100,11 @@ test:tpu --local_test_jobs=1
 common --experimental_repo_remote_exec
 
 # Inherit environmental variables that are used in testing.
-test --test_env=TPU_NUM_DEVICES --test_env=GPU_NUM_DEVICES --test_env=CPU_NUM_DEVICES --test_env=XRT_LOCAL_WORKER
+test --test_env=TPU_NUM_DEVICES --test_env=CPU_NUM_DEVICES --test_env=XRT_LOCAL_WORKER
 test --test_env=XRT_TPU_CONFIG --test_env=XRT_DEVICE_MAP --test_env=XRT_WORKERS --test_env=XRT_MESH_SERVICE_ADDRESS
 test --test_env=XRT_SHARD_WORLD_SIZE --test_env=XRT_MULTI_PROCESSING_DEVICE --test_env=XRT_HOST_ORDINAL --test_env=XRT_SHARD_ORDINAL
 test --test_env=XRT_START_LOCAL_SERVER --test_env=TPUVM_MODE --test_env=PJRT_DEVICE --test_env=PJRT_TPU_MAX_INFLIGHT_COMPUTATIONS
-test --test_env=PJRT_CPU_ASYNC_CLIENT --test_env=PJRT_GPU_ASYNC_CLIENT --test_env=TPU_LIBRARY_PATH --test_env=PJRT_DIST_SERVICE_ADDR
+test --test_env=PJRT_CPU_ASYNC_CLIENT --test_env=TPU_LIBRARY_PATH --test_env=PJRT_DIST_SERVICE_ADDR
 test --test_env=PJRT_LOCAL_PROCESS_RANK
 
 # This environmental variable is important for properly integrating with XLA.

BUILD

@@ -1,8 +1,3 @@
-load(
-    "@xla//xla/tsl/platform/default:cuda_build_defs.bzl",
-    "if_cuda_is_configured",
-)
-
 load("@python//:defs.bzl", "compile_pip_requirements")
 load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
 
@@ -41,9 +36,7 @@ cc_binary(
         "@torch//:libtorch",
         "@torch//:libtorch_cpu",
         "@torch//:libtorch_python",
-    ] + if_cuda_is_configured([
-        "@xla//xla/stream_executor:cuda_platform",
-    ]),
+    ],
 )
 
 test_suite(

WORKSPACE

@@ -56,8 +56,6 @@ http_archive(
     ],
     patch_tool = "patch",
     patches = [
-        "//openxla_patches:gpu_nvml.diff",
-        "//openxla_patches:gpu_race_condition.diff",
         "//openxla_patches:no_fortify.diff",
     ],
     strip_prefix = "xla-" + xla_hash,
@@ -142,16 +140,16 @@ load("@xla//:workspace0.bzl", "xla_workspace0")
 xla_workspace0()
 
 
+# Even though we don't support XLA:CUDA anymore, we still need to keep the
+# following. The reason being that `pjrt_computation_client_test` depends on
+# `@xla//xla/tools`, which calls:
+#
+# ```
+# load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")`
+# ```
 load(
     "@xla//third_party/gpus:cuda_configure.bzl",
     "cuda_configure",
 )
 
 cuda_configure(name = "local_config_cuda")
-
-load(
-    "@xla//third_party/nccl:nccl_configure.bzl",
-    "nccl_configure",
-)
-
-nccl_configure(name = "local_config_nccl")

openxla_patches/gpu_nvml.diff

@@ -2,10 +2,6 @@ load(
     "//bazel:rules_def.bzl",
     "ptxla_cc_test",
 )
-load(
-    "@xla//xla/tsl/platform/default:cuda_build_defs.bzl",
-    "if_cuda_is_configured",
-)
 
 licenses(["notice"])  # Apache 2.0
 
@@ -134,7 +130,6 @@ cc_library(
         "@xla//xla:shape_util",
         "@xla//xla/hlo/builder:xla_computation",
         "@xla//xla/pjrt:pjrt_client",
-        "@xla//xla/pjrt/c:pjrt_c_api_gpu_extension_hdrs",
         "@xla//xla/pjrt/c:pjrt_c_api_hdrs",
         "@xla//xla/pjrt/c:pjrt_c_api_wrapper_impl",
         "@xla//xla/pjrt:pjrt_c_api_client",
@@ -216,10 +211,9 @@ cc_library(
         "@torch//:headers",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:initialize",
+        "@xla//xla/pjrt/distributed:in_memory_key_value_store",
         "@xla//xla/pjrt:pjrt_c_api_client",
         "@xla//xla/pjrt:tfrt_cpu_pjrt_client",
-        "@xla//xla/pjrt/gpu:se_gpu_pjrt_client",
-        "@xla//xla/service:gpu_plugin",
     ],
 )
 
@@ -295,9 +289,7 @@ cc_library(
     deps = [
         "@xla//xla/backends/profiler/cpu:host_tracer",
         "@xla//xla/backends/profiler/cpu:metadata_collector",
-    ] + if_cuda_is_configured([
-        "@xla//xla/backends/profiler/gpu:device_tracer",
-    ]),
+    ],
     alwayslink = True,
 )