mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-28 03:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,438 Commits 508 Branches 263 Tags
Commit Graph

5 Commits

Author SHA1 Message Date
Benjamin Chetioui
875099b25d [Mosaic GPU] Enable the new transform inference pass in the warpgroup lowering. 
A couple of dummy transform inference rules needed to be added in order to
contend with parts of the lowering that do not use the dialect yet, along with
a transform inference rule for `memref.view`.

PiperOrigin-RevId: 738089782
 2025-03-18 11:51:43 -07:00
Benjamin Chetioui
ba2f7c9ad9 [Mosaic GPU] Add transform inference rule for mgpu.slice_smem. 
PiperOrigin-RevId: 737957778
 2025-03-18 04:53:54 -07:00
Benjamin Chetioui
9a686e0bf3 [Mosaic GPU] Add initial transform inference rules for vector.{load,store}. 
PiperOrigin-RevId: 737703568
 2025-03-17 12:08:07 -07:00
Benjamin Chetioui
d09df7c8ab [Mosaic GPU] Add transform inference rules for mgpu.async_{load,store}. 
PiperOrigin-RevId: 736795784
 2025-03-14 02:37:55 -07:00
Benjamin Chetioui
d028354abb [Mosaic GPU] Introduce an initial transform inference pass. 
For now, propagate transforms for `wgmma`. We do not handle `transpose` for
either operand yet.

The pass isn't called anywhere yet.

PiperOrigin-RevId: 736758754
 2025-03-13 23:22:59 -07:00