mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Change CLANG_CUDA_COMPILER_PATH set order. Add --config=cuda_clang to build.py

Browse Source 
Set `--action_env=CLANG_CUDA_COMPILER_PATH` after cuda_nvcc configuration
Add `--config=cuda_clang` when `--nouse_cuda_nvcc` flag set

PiperOrigin-RevId: 678873849
This commit is contained in:
jax authors 2024-09-25 15:39:09 -07:00
parent ce99c18a74
commit f1b3251bf9
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
build/build.py
View File

@ -284,11 +284,13 @@ def write_bazelrc(*, remote_build,
f.write("build --config=mkl_open_source_only\n")
if enable_cuda:
f.write("build --config=cuda\n")
if use_cuda_nvcc:
f.write("build --config=cuda_nvcc\n")
else:
f.write("build --config=cuda_clang\n")
f.write(f"build --action_env=CLANG_CUDA_COMPILER_PATH={clang_path}\n")
if not enable_nccl:
f.write("build --config=nonccl\n")
if use_cuda_nvcc:
f.write("build --config=cuda_nvcc\n")
if cuda_version:
f.write("build --repo_env HERMETIC_CUDA_VERSION=\"{cuda_version}\"\n"
.format(cuda_version=cuda_version))