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
rocm_jax/jaxlib/mosaic
History
Dimitar (Mitko) Asenov 3b305c6617 [Mosaic GPU] Infer layouts (transforms) on memrefs that directly feed into the dialect wgmma op. 
This change detects a situation where a gmem_memref is read via `async_load` and directly used in a wgmma. In such cases, we insert a cast before the load to add tile, transpose, and swizzle transformations.

PiperOrigin-RevId: 732618760
2025-03-02 03:17:13 -08:00
..
dialect
[Mosaic GPU] Infer layouts (transforms) on memrefs that directly feed into the dialect wgmma op.
2025-03-02 03:17:13 -08:00
gpu
Integrate LLVM at llvm/llvm-project@9d24f94379
2025-02-18 10:30:48 -08:00
python
Fix OSS build for the Mosaic GPU dialect
2025-01-02 15:55:03 +00:00
BUILD
Fix build dependencies
2025-02-26 08:38:31 -08:00
serde.cc
[mosaic] Extracted serialization pass traversal logic into a reusable function
2025-01-16 02:58:06 -08:00
serde.h
[mosaic] Extracted serialization pass traversal logic into a reusable function
2025-01-16 02:58:06 -08:00