Move tsl/cuda to xla/tsl/cuda

PiperOrigin-RevId: 610550833

Author: ddunl

jaxlib/BUILD

@@ -200,8 +200,8 @@ pybind_extension(
         # TODO(jieying): move to jaxlib after py_client_gpu is separated from py_client
         "@xla//xla/python:py_client_gpu",
         "@xla//xla/python:status_casters",
-        "@tsl//tsl/cuda:cublas",
-        "@tsl//tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cublas",
+        "@xla//xla/tsl/cuda:cudart",
     ],
 )
 

jaxlib/cuda/BUILD

@@ -14,13 +14,13 @@
 
 # NVIDIA CUDA kernels
 
-load("@rules_python//python:defs.bzl", "py_library")
 load(
     "//jaxlib:jax.bzl",
     "cuda_library",
     "if_cuda_is_configured",
     "pybind_extension",
 )
+load("@rules_python//python:defs.bzl", "py_library")
 
 licenses(["notice"])
 
@@ -56,9 +56,9 @@ cc_library(
     features = ["-use_header_modules"],
     deps = [
         ":cuda_vendor",
-        "@tsl//tsl/cuda:cupti",
-        "@tsl//tsl/cuda:cusolver",
-        "@tsl//tsl/cuda:cusparse",
+        "@xla//xla/tsl/cuda:cupti",
+        "@xla//xla/tsl/cuda:cusolver",
+        "@xla//xla/tsl/cuda:cusparse",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/memory",
@@ -81,8 +81,8 @@ cc_library(
         "//jaxlib:handle_pool",
         "//jaxlib:kernel_helpers",
         "@xla//xla/service:custom_call_status",
-        "@tsl//tsl/cuda:cublas",
-        "@tsl//tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cublas",
+        "@xla//xla/tsl/cuda:cudart",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",
@@ -117,7 +117,7 @@ pybind_extension(
         ":cublas_kernels",
         ":cuda_vendor",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cublas",
+        "@xla//xla/tsl/cuda:cublas",
         "@tsl//tsl/python/lib/core:numpy",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings:str_format",
@@ -135,8 +135,8 @@ cc_library(
         "//jaxlib:handle_pool",
         "//jaxlib:kernel_helpers",
         "@xla//xla/service:custom_call_status",
-        "@tsl//tsl/cuda:cudart",
-        "@tsl//tsl/cuda:cudnn",
+        "@xla//xla/tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cudnn",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
@@ -174,7 +174,7 @@ cc_library(
         "//jaxlib:handle_pool",
         "//jaxlib:kernel_helpers",
         "@xla//xla/service:custom_call_status",
-        "@tsl//tsl/cuda:cusolver",
+        "@xla//xla/tsl/cuda:cusolver",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
@@ -203,8 +203,8 @@ pybind_extension(
         ":cuda_vendor",
         ":cusolver_kernels",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cudart",
-        "@tsl//tsl/cuda:cusolver",
+        "@xla//xla/tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cusolver",
         "@tsl//tsl/python/lib/core:numpy",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings:str_format",
@@ -223,8 +223,8 @@ cc_library(
         "//jaxlib:handle_pool",
         "//jaxlib:kernel_helpers",
         "@xla//xla/service:custom_call_status",
-        "@tsl//tsl/cuda:cudart",
-        "@tsl//tsl/cuda:cusparse",
+        "@xla//xla/tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cusparse",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
@@ -253,8 +253,8 @@ pybind_extension(
         ":cuda_vendor",
         ":cusparse_kernels",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cudart",
-        "@tsl//tsl/cuda:cusparse",
+        "@xla//xla/tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cusparse",
         "@tsl//tsl/python/lib/core:numpy",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
@@ -316,7 +316,7 @@ pybind_extension(
         ":cuda_lu_pivot_kernels_impl",
         ":cuda_vendor",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cudart",
         "@local_config_cuda//cuda:cuda_headers",
         "@nanobind",
     ],
@@ -366,7 +366,7 @@ pybind_extension(
         ":cuda_gpu_kernel_helpers",
         ":cuda_prng_kernels",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cudart",
         "@local_config_cuda//cuda:cuda_headers",
         "@nanobind",
     ],
@@ -400,7 +400,7 @@ cc_library(
         "//jaxlib/gpu:triton_cc_proto",
         "@xla//xla/service:custom_call_status",
         "@xla//xla/stream_executor/gpu:asm_compiler",
-        "@tsl//tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cudart",
         "@tsl//tsl/platform:env",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
@@ -472,13 +472,13 @@ cc_library(
         ":cuda_vendor",
         "//jaxlib:absl_status_casters",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cublas",
-        "@tsl//tsl/cuda:cudart",
-        "@tsl//tsl/cuda:cudnn",
-        "@tsl//tsl/cuda:cufft",
-        "@tsl//tsl/cuda:cupti",
-        "@tsl//tsl/cuda:cusolver",
-        "@tsl//tsl/cuda:cusparse",
+        "@xla//xla/tsl/cuda:cublas",
+        "@xla//xla/tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cudnn",
+        "@xla//xla/tsl/cuda:cufft",
+        "@xla//xla/tsl/cuda:cupti",
+        "@xla//xla/tsl/cuda:cusolver",
+        "@xla//xla/tsl/cuda:cusparse",
     ],
 )
 
@@ -509,13 +509,13 @@ pybind_extension(
         ":versions_helpers",
         "//jaxlib:absl_status_casters",
         "//jaxlib:kernel_nanobind_helpers",
-        "@tsl//tsl/cuda:cublas",
-        "@tsl//tsl/cuda:cudart",
-        "@tsl//tsl/cuda:cudnn",
-        "@tsl//tsl/cuda:cufft",
-        "@tsl//tsl/cuda:cupti",
-        "@tsl//tsl/cuda:cusolver",
-        "@tsl//tsl/cuda:cusparse",
+        "@xla//xla/tsl/cuda:cublas",
+        "@xla//xla/tsl/cuda:cudart",
+        "@xla//xla/tsl/cuda:cudnn",
+        "@xla//xla/tsl/cuda:cufft",
+        "@xla//xla/tsl/cuda:cupti",
+        "@xla//xla/tsl/cuda:cusolver",
+        "@xla//xla/tsl/cuda:cusparse",
         "@com_google_absl//absl/status:statusor",
         "@nanobind",
     ],