jax
2c382f86 - [Mosaic GPU] Use TiledLayout for TMEM layouts

Commit

204 days ago

[Mosaic GPU] Use TiledLayout for TMEM layouts TMEM is actually remarkably similar to the register file: it's partitioned into 4 banks that can't communicate (just like no communication is possible between the register files in each of the 4 quadrants in an SM), and each one of those banks has 32 lanes. Overall this approach seems to work quite well and should hopefully let us easily implement new layouts such as those for MMA with m=64. PiperOrigin-RevId: 774816973

References

#29603 - [Mosaic GPU] Use TiledLayout for TMEM layouts

#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

#581 - Fix/pallas tests shared memory

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

#585 - update a test for checking zero ROCm GPU event