mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update --config=cuda to add direct dependencies on CUDA libraries both for bazel build and bazel test phases.

Browse Source 
With this configuration the same cache is used both for `bazel build` and `bazel test` commands (provided the same target is specified).

Add `--config=no_cuda_libs` for building targets with CUDA libraries from stubs.

PiperOrigin-RevId: 720334587
This commit is contained in:
jax authors 2025-01-27 15:45:43 -08:00
parent ccd78b5b62
commit 727d0367a4
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.bazelrc
View File

@ -124,9 +124,10 @@ build:cuda --@local_config_cuda//:enable_cuda
# Default hermetic CUDA and CUDNN versions.
build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.3.2"
build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.1.1"
build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
# This flag is needed to include CUDA libraries for bazel tests.
test:cuda --@local_config_cuda//cuda:include_cuda_libs=true
# This config is used for building targets with CUDA libraries from stubs.
build:cuda_libraries_from_stubs --@local_config_cuda//cuda:include_cuda_libs=false
# Force the linker to set RPATH, not RUNPATH. When resolving dynamic libraries,
# ld.so prefers in order: RPATH, LD_LIBRARY_PATH, RUNPATH. JAX sets RPATH to

1
build/build.py
View File

@ -532,6 +532,7 @@ async def main():
if "cuda" in args.wheels:
wheel_build_command_base.append("--config=cuda")
wheel_build_command_base.append("--config=cuda_libraries_from_stubs")
if args.use_clang:
wheel_build_command_base.append(
f"--action_env=CLANG_CUDA_COMPILER_PATH=\"{clang_path}\""