jax
cbd0844e - [Mosaic GPU] Fold non-gather indices into the gather TMA column index

Commit

126 days ago

[Mosaic GPU] Fold non-gather indices into the gather TMA column index The offset is shared by all gather indices so this method reduces the overall amount of arithmetic instructions we emit. For as long as we don't add support for inputs that are more than 2D this essentially undoes any of the ops inserted by the tiling transforms (and so the only downside of tiling is that we might use more TMA instructions, but not more arithmetic). This implementation is simpler too. PiperOrigin-RevId: 796798377

References

#31122 - [Mosaic GPU] Fold non-gather indices into the gather TMA column index

#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