mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Mosaic GPU] Add layout inference for arith.Ext{F,SI,UI}Op and arith.Trunc{F,I}Op.

Browse Source 
PiperOrigin-RevId: 730851596
This commit is contained in:
Benjamin Chetioui 2025-02-25 05:58:58 -08:00 committed by jax authors
parent 6f966397e0
commit 5312b5e35a
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/experimental/mosaic/gpu/layout_inference.py
View File

@ -243,6 +243,7 @@ for op in [
arith.AndIOp,
arith.CmpFOp,
arith.CmpIOp,
arith.ExtFOp, arith.ExtSIOp, arith.ExtUIOp,
arith.MaximumFOp,
arith.MaxUIOp, arith.MaxSIOp,
arith.MinimumFOp,
@ -252,6 +253,7 @@ for op in [
arith.FloorDivSIOp, arith.DivUIOp, arith.DivFOp,
arith.RemUIOp, arith.RemSIOp, arith.RemFOp,
arith.SubIOp, arith.SubFOp,
arith.TruncFOp, arith.TruncIOp,
arith.XOrIOp,
vector.LoadOp,
vector.StoreOp,