Repository: jax-ml/jax
Pull request #27099 (2025-03-12 19:12:25)

[Mosaic GPU] Add transform inference rules for `mgpu.async_{load,store}`.
[Mosaic GPU] Add transform inference rules for `mgpu.async_{load,store}`.



Patch:
diff --git a/jax/experimental/mosaic/gpu/transform_inference.py b/jax/experimental/mosaic/gpu/transform_inference.py
index 5a275af66..65936f943 100644
--- a/jax/experimental/mosaic/gpu/transform_inference.py
+++ b/jax/experimental/mosaic/gpu/transform_inference.py
@@ -98,6 +98,18 @@ def infer_wgmma_transforms(op: mgpu.WGMMAOp) -> OptionalTransforms:
   return [b_transforms], []
 
 
+@partial(_add_transform_inference_rule, mgpu.AsyncStoreOp)
+def _infer_async_store_transforms(op: mgpu.AsyncStoreOp) -> OptionalTransforms:
+  in_transforms = inference_utils.value_transforms(op.source)
+  return None if in_transforms is None else ([in_transforms], [])
+
+
+@partial(_add_transform_inference_rule, mgpu.AsyncLoadOp)
+def _infer_async_load_transforms(op: mgpu.AsyncLoadOp) -> OptionalTransforms:
+  in_transforms = inference_utils.value_transforms(op.destination)
+  return None if in_transforms is None else ([in_transforms], [])
+
+
 def _should_have_transforms(op: ir.OpView) -> bool:
   """Returns 'True' if the operation should be assigned in/out transforms."""
   return any(
@@ -111,7 +123,10 @@ def _should_have_transforms(op: ir.OpView) -> bool:
 def infer_transforms(module: ir.Module):
   """Infers transforms for the given module.
 
-  Transforms are to smem refs the what layouts are to vectors.
+  Transforms are to memrefs what layouts are to vectors. More specifically,
+  transforms describe mappings between SMEM refs and GMEM refs, and are
+  determined based on how SMEM refs are used. For that reason, we always
+  annotate and apply memrefs on SMEM refs.
 
   The pass is meant to be called on a module where layouts have been fully
   specified. We error out if two distinct sets of transforms are competing to


Test patch:
diff --git a/tests/mosaic/gpu_transform_inference_test.py b/tests/mosaic/gpu_transform_inference_test.py
index 2bae3fcd0..65c851dda 100644
--- a/tests/mosaic/gpu_transform_inference_test.py
+++ b/tests/mosaic/gpu_transform_inference_test.py
@@ -23,6 +23,7 @@ from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.interpreters import mlir as mlir_interpreter
 from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import arith
 from jax._src.lib.mlir.dialects import func
 import jax.experimental.mosaic.gpu as mgpu
 from jax.experimental.mosaic.gpu import inference_utils
@@ -92,6 +93,75 @@ class TransformInferenceTest(parameterized.TestCase):
     )
     self.assertEmpty(inference_utils.out_transforms(wgmma_op))
 
+  def test_infer_transforms_for_async_load_derives_from_destination(self):
+    async_load_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(gmem_ref, smem_ref, barrier):
+      nonlocal async_load_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      async_load_op = mgpu.dialect.AsyncLoadOp(
+          source=gmem_ref,
+          destination=smem_ref,
+          barrier=barrier,
+          indices=[zero, zero],
+          slice_lengths=shape,
+          collective=ir.ArrayAttr.get([]),
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      gmem_ty = ir.MemRefType.get(shape, elt_ty)
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      barrier_ty = ir.Type.parse("!mosaic_gpu.barrier")
+      f = func.FuncOp.from_py_func(gmem_ty, smem_ty, barrier_ty)(body).func_op
+
+    transforms = ir.ArrayAttr.get(
+        [mgpu.dialect.TransposeTransformAttr.get((1, 0))]
+    )
+    f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    mgpu.infer_transforms(self.module)
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(async_load_op), [transforms]
+    )
+    self.assertEmpty(inference_utils.out_transforms(async_load_op))
+
+  def test_infer_transforms_for_async_store_op_derives_from_source(self):
+    async_store_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(gmem_ref, smem_ref):
+      nonlocal async_store_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      async_store_op = mgpu.dialect.AsyncStoreOp(
+          source=smem_ref,
+          destination=gmem_ref,
+          indices=[zero, zero],
+          slice_lengths=shape,
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      gmem_ty = ir.MemRefType.get(shape, elt_ty)
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      f = func.FuncOp.from_py_func(gmem_ty, smem_ty)(body).func_op
+
+    transforms = ir.ArrayAttr.get(
+        [mgpu.dialect.TransposeTransformAttr.get((1, 0))]
+    )
+    f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    mgpu.infer_transforms(self.module)
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(async_store_op), [transforms]
+    )
+    self.assertEmpty(inference_utils.out_transforms(async_store_op))
+
 
 if __name__ == "__main__":
   parameterized.absltest.main(testLoader=jtu.JaxTestLoader())