[Mosaic GPU] Add initial transform inference rules for vector.{load,store}.

PiperOrigin-RevId: 737703568

Author: bchetioui

jax/experimental/mosaic/gpu/transform_inference.py

@@ -25,8 +25,12 @@
 
 from jax._src.lib import mosaic_gpu_dialect as mgpu
 from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import arith
+from jax._src.lib.mlir.dialects import vector
 
+from . import fragmented_array as fa
 from . import inference_utils
+from . import layouts as layouts_lib
 from . import utils
 
 # mypy: ignore-errors
@@ -40,6 +44,7 @@ def _add_transform_inference_rule(
     op: type[ir.OpView], rule: TransformInferenceRule
 ):
   _transform_inference_rules[op.OPERATION_NAME] = rule  # pytype: disable=attribute-error
+  return rule
 
 
 def _set_transform_attributes(
@@ -110,6 +115,60 @@ def _infer_async_load_transforms(op: mgpu.AsyncLoadOp) -> OptionalTransforms:
   return None if in_transforms is None else ([in_transforms], [])
 
 
+@partial(_add_transform_inference_rule, vector.LoadOp)
+@partial(_add_transform_inference_rule, vector.StoreOp)
+def _infer_vector_load_store_transforms(
+    op: vector.LoadOp | vector.StoreOp,
+) -> OptionalTransforms:
+  for i in op.indices:
+    index_defining_op = i.owner.opview
+    if (
+        not isinstance(index_defining_op, arith.ConstantOp)
+        or index_defining_op.literal_value != 0
+    ):
+      # TODO(bchetioui): handle slicing.
+      raise NotImplementedError(
+          f"Only constants with value 0 are supported as indices for {op}"
+      )
+
+  if isinstance(op, vector.LoadOp):
+    [layout_attr] = inference_utils.out_layouts(op)
+  else:
+    assert isinstance(op, vector.StoreOp)
+    [layout_attr] = inference_utils.in_layouts(op)
+
+  layout = layouts_lib.from_layout_attr(layout_attr)
+  transforms = inference_utils.value_transforms(op.base)
+
+  if layout == fa.WGMMA_LAYOUT:
+    layout_transforms = infer_transforms_for_wgmma_ref(
+        ir.MemRefType(op.base.type)
+    )
+  elif (isinstance(layout, fa.WGStridedFragLayout) or
+        isinstance(layout, fa.WGSplatFragLayout)):
+    layout_transforms = None
+  else:
+    raise NotImplementedError(
+        f"Got layout {layout} which is not yet supported"
+    )
+
+  if transforms is not None and layout_transforms is not None:
+    if transforms != layout_transforms:
+      raise NotImplementedError(
+          f"Conflicting transforms for {op.base} in {op}: "
+          f"{transforms} != {layout_transforms}."
+      )
+    return [transforms], []
+
+  if transforms is not None:
+    return [transforms], []
+
+  if layout_transforms is not None:
+    return [layout_transforms], []
+
+  return None
+
+
 def _should_have_transforms(op: ir.OpView) -> bool:
   """Returns 'True' if the operation should be assigned in/out transforms."""
   return any(

tests/mosaic/gpu_transform_inference_test.py

@@ -25,8 +25,11 @@
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import arith
 from jax._src.lib.mlir.dialects import func
+from jax._src.lib.mlir.dialects import vector
 import jax.experimental.mosaic.gpu as mgpu
+from jax.experimental.mosaic.gpu import fragmented_array as fa
 from jax.experimental.mosaic.gpu import inference_utils
+from jax.experimental.mosaic.gpu import layouts as layouts_lib
 import numpy as np
 
 
@@ -162,6 +165,187 @@ def body(gmem_ref, smem_ref):
     )
     self.assertEmpty(inference_utils.out_transforms(async_store_op))
 
+  def test_infer_transforms_for_vector_load_op_derives_from_destination(self):
+    vector_load_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(smem_ref):
+      nonlocal vector_load_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      vector_load_op = vector.LoadOp(
+          ir.VectorType.get(shape, elt_ty), smem_ref, [zero] * len(shape)
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      func.FuncOp.from_py_func(smem_ty)(body)
+
+    vector_load_op.attributes["out_layouts"] = ir.ArrayAttr.get(
+        [layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)]
+    )
+
+    mgpu.infer_transforms(self.module)
+
+    expected_transforms = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((8, 64)),
+        mgpu.dialect.SwizzleTransformAttr.get(128),
+    ])
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(vector_load_op), [expected_transforms]
+    )
+    self.assertEmpty(inference_utils.out_transforms(vector_load_op))
+
+  def test_infer_transforms_for_vector_load_op_derives_from_source(self):
+    vector_load_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(smem_ref):
+      nonlocal vector_load_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      vector_load_op = vector.LoadOp(
+          ir.VectorType.get(shape, elt_ty), smem_ref, [zero] * len(shape)
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      f = func.FuncOp.from_py_func(smem_ty)(body).func_op
+
+    vector_load_op.attributes["out_layouts"] = ir.ArrayAttr.get(
+        [layouts_lib.to_layout_attr(fa.WGStridedFragLayout(shape, vec_size=4))]
+    )
+    transforms = ir.ArrayAttr.get([mgpu.dialect.TileTransformAttr.get((8, 64))])
+    f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    mgpu.infer_transforms(self.module)
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(vector_load_op), [transforms]
+    )
+    self.assertEmpty(inference_utils.out_transforms(vector_load_op))
+
+  def test_infer_transforms_for_vector_load_op_raises_on_mismatches(self):
+    vector_load_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(smem_ref):
+      nonlocal vector_load_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      vector_load_op = vector.LoadOp(
+          ir.VectorType.get(shape, elt_ty), smem_ref, [zero] * len(shape)
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      f = func.FuncOp.from_py_func(smem_ty)(body).func_op
+
+    vector_load_op.attributes["out_layouts"] = ir.ArrayAttr.get(
+        [layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)]
+    )
+    transforms = ir.ArrayAttr.get([mgpu.dialect.TileTransformAttr.get((8, 64))])
+    f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    with self.assertRaisesRegex(NotImplementedError, "Conflicting transforms"):
+      mgpu.infer_transforms(self.module)
+
+  def test_infer_transforms_for_vector_store_op_derives_from_destination(self):
+    vector_store_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(smem_ref, value_to_store):
+      nonlocal vector_store_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      vector_store_op = vector.StoreOp(
+          value_to_store, smem_ref, [zero] * len(shape)
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      value_ty = ir.VectorType.get(shape, elt_ty)
+      func.FuncOp.from_py_func(smem_ty, value_ty)(body)
+
+    vector_store_op.attributes["in_layouts"] = ir.ArrayAttr.get(
+        [layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)]
+    )
+
+    mgpu.infer_transforms(self.module)
+
+    expected_transforms = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((8, 64)),
+        mgpu.dialect.SwizzleTransformAttr.get(128),
+    ])
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(vector_store_op), [expected_transforms]
+    )
+    self.assertEmpty(inference_utils.out_transforms(vector_store_op))
+
+  def test_infer_transforms_for_vector_store_op_derives_from_source(self):
+    vector_store_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(smem_ref, value_to_store):
+      nonlocal vector_store_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      vector_store_op = vector.StoreOp(
+          value_to_store, smem_ref, [zero] * len(shape)
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      value_ty = ir.VectorType.get(shape, elt_ty)
+      f = func.FuncOp.from_py_func(smem_ty, value_ty)(body).func_op
+
+    vector_store_op.attributes["in_layouts"] = ir.ArrayAttr.get(
+        [layouts_lib.to_layout_attr(fa.WGStridedFragLayout(shape, vec_size=4))]
+    )
+    transforms = ir.ArrayAttr.get([mgpu.dialect.TileTransformAttr.get((8, 64))])
+    f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    mgpu.infer_transforms(self.module)
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(vector_store_op), [transforms]
+    )
+    self.assertEmpty(inference_utils.out_transforms(vector_store_op))
+
+  def test_infer_transforms_for_vector_store_op_raises_on_mismatches(self):
+    vector_store_op = None
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    def body(smem_ref, value_to_store):
+      nonlocal vector_store_op
+      zero = arith.constant(ir.IntegerType.get_signless(32), 0)
+      vector_store_op = vector.StoreOp(
+          value_to_store, smem_ref, [zero] * len(shape)
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      smem_ty = ir.MemRefType.get(shape, elt_ty, memory_space=smem)
+      value_ty = ir.VectorType.get(shape, elt_ty)
+      f = func.FuncOp.from_py_func(smem_ty, value_ty)(body).func_op
+
+    vector_store_op.attributes["in_layouts"] = ir.ArrayAttr.get(
+        [layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)]
+    )
+    transforms = ir.ArrayAttr.get([mgpu.dialect.TileTransformAttr.get((8, 64))])
+    f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    with self.assertRaisesRegex(NotImplementedError, "Conflicting transforms"):
+      mgpu.infer_transforms(self.module)
+
 
 if __name__ == "__main__":
   parameterized.absltest.main(testLoader=jtu.JaxTestLoader())