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

Include SASS/PTX for Hopper GPUs.

Browse Source
This commit is contained in:
Peter Hawkins 2023-06-05 09:42:12 -04:00
parent 025e348323
commit cb33fdf3f7
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.bazelrc
View File

@ -66,7 +66,7 @@ build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1
build:cuda --repo_env TF_NEED_CUDA=1
# "sm" means we emit only cubin, which is forward compatible within a GPU generation.
# "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="sm_52,sm_60,sm_70,compute_80"
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="sm_52,sm_60,sm_70,sm_80,compute_90"
build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
build:cuda --@local_config_cuda//:enable_cuda
build:cuda --@xla//xla/python:enable_gpu=true

5
CHANGELOG.md
View File

@ -22,6 +22,11 @@ Remember to align the itemized text with the first line of an item within a list
## jaxlib 0.4.12
* Changes
* Include PTX/SASS for Hopper (SM version 9.0+) GPUs. Previous
versions of jaxlib should work on Hopper but would have a long
JIT-compilation delay the first time a JAX operation was executed.
## jax 0.4.11 (May 31, 2023)
* Deprecations