Remove CUDA rpaths from jaxlib build.

These are also set in the TSL build rules as part of the CUDA stub libraries, which these libraries depend on, so these copies of the rpath settings are redundant. PiperOrigin-RevId: 716844265
jax-ml · Jan 18, 2025 · 034e967 · 034e967
1 parent c7f8d17
commit 034e967
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 58 deletions.
diff --git a/jaxlib/cuda/BUILD b/jaxlib/cuda/BUILD
@@ -133,13 +133,6 @@ nanobind_extension(
         "-fno-strict-aliasing",
     ],
     features = ["-use_header_modules"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_runtime/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cublas/lib",
-        ],
-        "//conditions:default": [],
-    }),
     module_name = "_blas",
     deps = [
         ":cublas_kernels",
@@ -273,14 +266,6 @@ nanobind_extension(
         "-fno-strict-aliasing",
     ],
     features = ["-use_header_modules"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_runtime/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cusolver/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cublas/lib",
-        ],
-        "//conditions:default": [],
-    }),
     module_name = "_solver",
     deps = [
         ":cuda_gpu_kernel_helpers",
@@ -329,13 +314,6 @@ nanobind_extension(
         "-fno-strict-aliasing",
     ],
     features = ["-use_header_modules"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_runtime/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cusparse/lib",
-        ],
-        "//conditions:default": [],
-    }),
     module_name = "_sparse",
     deps = [
         ":cuda_gpu_kernel_helpers",
@@ -505,12 +483,6 @@ nanobind_extension(
         "-fno-strict-aliasing",
     ],
     features = ["-use_header_modules"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_runtime/lib",
-        ],
-        "//conditions:default": [],
-    }),
     module_name = "_hybrid",
     deps = [
         ":cuda_gpu_kernel_helpers",
@@ -597,12 +569,6 @@ nanobind_extension(
         "-fno-strict-aliasing",
     ],
     features = ["-use_header_modules"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_runtime/lib",
-        ],
-        "//conditions:default": [],
-    }),
     module_name = "_triton",
     deps = [
         ":cuda_gpu_kernel_helpers",
@@ -648,18 +614,6 @@ nanobind_extension(
         "-fno-strict-aliasing",
     ],
     features = ["-use_header_modules"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_cupti/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cuda_runtime/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cublas/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cufft/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cudnn/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cusolver/lib",
-            "-Wl,-rpath,$$ORIGIN/../../nvidia/cusparse/lib",
-        ],
-        "//conditions:default": [],
-    }),
     module_name = "_versions",
     deps = [
         ":cuda_gpu_kernel_helpers",

diff --git a/jaxlib/mosaic/gpu/BUILD b/jaxlib/mosaic/gpu/BUILD
@@ -181,12 +181,6 @@ nanobind_extension(
         "-fexceptions",
         "-fno-strict-aliasing",
     ],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../../nvidia/cuda_runtime/lib",
-        ],
-        "//conditions:default": [],
-    }),
     deps = [
         ":mlir_capi",
         "//jaxlib:kernel_nanobind_helpers",
@@ -204,12 +198,6 @@ cc_binary(
     name = "libmosaic_gpu_runtime.so",
     srcs = ["runtime.cc"],
     copts = ["-fvisibility=default"],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../../nvidia/cuda_runtime/lib",
-        ],
-        "//conditions:default": [],
-    }),
     linkshared = 1,
     tags = [
         "manual",