jax
967489d8 - [Pallas:MGPU] Make TMEM reads and writes explicitly asynchronous

Commit

190 days ago

[Pallas:MGPU] Make TMEM reads and writes explicitly asynchronous We used to pretend that they are synchronous operations, same as for SMEM, but this would be too expensive to guarantee. We now explicitly disallow explicit loads and store to TMEM refs and instead require the user to use * `plgpu.async_tmem_load` + `plgpu.wait_tmem_load` to await it * `plgpu.async_tmem_store` + `plgpu.commit_tmem` to await it PiperOrigin-RevId: 778021794

References

#29886 - [Pallas:MGPU] Make TMEM reads and writes explicitly asynchronous

#31381 - Remove old ROCm build code

#31720 - Fix ann_test.py numerical bug in target reshape

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

#579 - Create rocm-test-requirements.txt

#580 - GESVDJ support for ROCm GPUs in JAX

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency