Repository: jax-ml/jax
Pull request #27097 (2025-03-12 18:27:33)

[Mosaic GPU] Introduce an initial transform inference pass.
[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.



Patch:
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index c485fd3a7..d004c7deb 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -45,6 +45,10 @@ from .layout_inference import (
     infer_layout as infer_layout,
 )
 
+from .transform_inference import (
+    infer_transforms as infer_transforms,
+)
+
 from .fragmented_array import (
     FragmentedArray as FragmentedArray,
     FragmentedLayout as FragmentedLayout,
diff --git a/jax/experimental/mosaic/gpu/inference_utils.py b/jax/experimental/mosaic/gpu/inference_utils.py
index cad3bfc11..87ef3a791 100644
--- a/jax/experimental/mosaic/gpu/inference_utils.py
+++ b/jax/experimental/mosaic/gpu/inference_utils.py
@@ -45,6 +45,28 @@ def out_layouts(op: MlirOperation) -> Sequence[ir.Attribute]:
   return op.attributes["out_layouts"]  # type: ignore
 
 
+def in_transforms(op: MlirOperation) -> Sequence[ir.Attribute]:
+  """Returns the in_transforms attribute of the given operation.
+
+  Raises:
+    ValueError: If the operation does not have an in_transforms attribute.
+  """
+  if "in_transforms" not in op.attributes:
+    raise ValueError(f"{op} does not have an in_transforms attribute.")
+  return op.attributes["in_transforms"]  # type: ignore
+
+
+def out_transforms(op: MlirOperation) -> Sequence[ir.Attribute]:
+  """Returns the out_transforms attribute of the given operation.
+
+  Raises:
+    ValueError: If the operation does not have an out_transforms attribute.
+  """
+  if "out_transforms" not in op.attributes:
+    raise ValueError(f"{op} does not have an out_transforms attribute.")
+  return op.attributes["out_transforms"]  # type: ignore
+
+
 def should_have_layout(op: MlirOperation) -> bool:
   """Returns 'true' if the operation should be assigned a layout."""
 
@@ -64,6 +86,33 @@ def has_any_layout_set(op: MlirOperation) -> bool:
   return has_in_layouts_set(op) or has_out_layouts_set(op)
 
 
+def has_in_transforms_set(op: MlirOperation) -> bool:
+  return "in_transforms" in op.attributes
+
+
+def has_out_transforms_set(op: MlirOperation) -> bool:
+  return "out_transforms" in op.attributes
+
+
+def _in_attr_for_operand(
+    op: MlirOperation,
+    operand: ir.Value,
+    attr_name: str,
+) -> ir.Attribute | None:
+  if attr_name == "in_layouts":
+    predicate = lambda v: ir.VectorType.isinstance(v.type)
+  elif attr_name == "in_transforms":
+    predicate = lambda v: ir.MemRefType.isinstance(v.type)
+  else:
+    raise ValueError(f"Unknown attribute: {attr_name}")
+
+  operand_number = [o for o in op.operands if predicate(o)].index(operand)
+
+  if attr_name not in op.attributes:
+    return None
+  return op.attributes[attr_name][operand_number]  # type: ignore
+
+
 def in_layout_for_operand(
     op: MlirOperation,
     operand: ir.Value,
@@ -77,51 +126,97 @@ def in_layout_for_operand(
   if not ir.VectorType.isinstance(operand.type):
     raise ValueError(f"{operand} is not a vector.")
 
-  operand_number = [
-      o for o in op.operands if ir.VectorType.isinstance(o.type)
-  ].index(operand)
-
-  if not has_in_layouts_set(op):
-    return None
-
-  return in_layouts(op)[operand_number]
+  return _in_attr_for_operand(op, operand, "in_layouts")
 
 
-def value_layout(value: ir.Value) -> ir.Attribute | None:
-  """Returns the layout for a given value as defined by its owner.
+def in_transforms_for_operand(
+    op: MlirOperation,
+    operand: ir.Value,
+) -> ir.Attribute | None:
+  """Returns the transforms for the operand in the given operation if it is set.
 
   Raises:
-    ValueError: If `result` is not a Vector.
+    ValueError: If `operand` is not an operand of `op`, or if `operand` is not a
+      memref.
   """
-  if not ir.VectorType.isinstance(value.type):
-    raise ValueError(f"{value} is not a vector.")
+  if not ir.MemRefType.isinstance(operand.type):
+    raise ValueError(f"{operand} is not a memref.")
+
+  return _in_attr_for_operand(op, operand, "in_transforms")
+
+
+def is_transformable_smem_memref(v: ir.Value) -> bool:
+  """Whether the value is a memref in SMEM on which transforms should be applied."""
+  barrier_ty = ir.Type.parse("!mosaic_gpu.barrier")
+  smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+  return (
+      ir.MemRefType.isinstance(v.type)
+      # barriers have no business being transformed
+      and v.type.element_type != barrier_ty  # pylint: disable=attribute-error
+      and v.type.memory_space is not None  # pylint: disable=attribute-error
+      and v.type.memory_space == smem  # pylint: disable=attribute-error
+  )
+
+
+def _value_attr(value: ir.Value, attr_type: str) -> ir.Attribute | None:
+  if attr_type == "layouts":
+    predicate = lambda v: ir.VectorType.isinstance(v.type)
+  elif attr_type == "transforms":
+    predicate = is_transformable_smem_memref
+  else:
+    raise ValueError(f"Unknown attribute: {attr_type}")
+
+  in_attr_type = "in_" + attr_type
+  out_attr_type = "out_" + attr_type
 
   owner = value.owner
   if isinstance(owner, ir.Operation):
-    if not has_out_layouts_set(owner):
+    if out_attr_type not in owner.attributes:
       return None
-    value_result_number = [
-        r for r in owner.results if ir.VectorType.isinstance(r.type)
-    ].index(value)
-    return out_layouts(owner)[value_result_number]
+    value_result_number = [r for r in owner.results if predicate(r)].index(
+        value
+    )
+    return owner.attributes[out_attr_type][value_result_number]  # type: ignore
 
   # Block case, useful when attempting to derive layouts for ops
   # depending on function parameters, or loop block arguments.
   if isinstance(owner, ir.Block):
     owner_op = owner.owner
     block = cast(ir.Block, owner)
-    if not has_in_layouts_set(owner_op):
+    if in_attr_type not in owner_op.attributes:
       return None
-    value_arg_number = [
-        r for r in block.arguments if ir.VectorType.isinstance(r.type)
-    ].index(value)
-    return in_layouts(owner_op)[value_arg_number]
+    value_arg_number = [r for r in block.arguments if predicate(r)].index(value)
+    return owner_op.attributes[in_attr_type][value_arg_number]  # type: ignore
 
   raise NotImplementedError(
       f"{owner} is not a function block nor an operation."
   )
 
 
+def value_layout(value: ir.Value) -> ir.Attribute | None:
+  """Returns the layout for a given value as defined by its owner.
+
+  Raises:
+    ValueError: If `result` is not a Vector.
+  """
+  if not ir.VectorType.isinstance(value.type):
+    raise ValueError(f"{value} is not a vector.")
+
+  return _value_attr(value, "layouts")
+
+
+def value_transforms(value: ir.Value) -> ir.Attribute | None:
+  """Returns the transforms for a given value as defined by its owner.
+
+  Raises:
+    ValueError: If `result` is not a memref.
+  """
+  if not ir.MemRefType.isinstance(value.type):
+    raise ValueError(f"{value} is not a memref.")
+
+  return _value_attr(value, "transforms")
+
+
 class TraversalOrder(enum.Enum):
   """Traversal orders with respect to the data flow for IR."""
 
diff --git a/jax/experimental/mosaic/gpu/transform_inference.py b/jax/experimental/mosaic/gpu/transform_inference.py
new file mode 100644
index 000000000..5a275af66
--- /dev/null
+++ b/jax/experimental/mosaic/gpu/transform_inference.py
@@ -0,0 +1,145 @@
+# Copyright 2025 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transform inference pass for the MLIR Mosaic GPU dialect.
+
+The transform inference pass is meant to run on IR that has already been
+annotated with layouts (see `layout_inference.py` for the relevant pass).
+"""
+
+from collections.abc import Callable
+from functools import partial
+import itertools
+from typing import cast
+
+from jax._src.lib import mosaic_gpu_dialect as mgpu
+from jax._src.lib.mlir import ir
+
+from . import inference_utils
+from . import utils
+
+# mypy: ignore-errors
+
+OptionalTransforms = tuple[list[ir.Attribute], list[ir.Attribute]] | None
+TransformInferenceRule = Callable[[ir.OpView], OptionalTransforms]
+_transform_inference_rules: dict[str, TransformInferenceRule] = {}
+
+
+def _add_transform_inference_rule(
+    op: type[ir.OpView], rule: TransformInferenceRule
+):
+  _transform_inference_rules[op.OPERATION_NAME] = rule  # pytype: disable=attribute-error
+
+
+def _set_transform_attributes(
+    op: ir.OpView,
+    in_transforms: list[ir.Attribute],
+    out_transforms: list[ir.Attribute],
+):
+  op.attributes["in_transforms"] = ir.ArrayAttr.get(in_transforms)
+  op.attributes["out_transforms"] = ir.ArrayAttr.get(out_transforms)
+
+
+def infer_transforms_for_wgmma_ref(ref_ty: ir.MemRefType) -> ir.ArrayAttr:
+  if len(ref_ty.shape) != 2:
+    raise ValueError(f"Expected a 2D memref, got {ref_ty}")
+
+  element_bytewidth = utils.bytewidth(ref_ty.element_type)
+  strides, _ = ref_ty.get_strides_and_offset()
+
+  if strides[0] < strides[1]:
+    raise NotImplementedError("Transpositions aren't handled yet.")
+
+  minor_dim = ref_ty.shape[1]
+  major_tiling = 8
+
+  # Try tiling with all swizzling modes starting from the largest one.
+  for swizzle in [
+      mgpu.SwizzlingMode.k128ByteSwizzle,
+      mgpu.SwizzlingMode.k64ByteSwizzle,
+      mgpu.SwizzlingMode.k32ByteSwizzle,
+      mgpu.SwizzlingMode.kNoSwizzle,
+  ]:
+    swizzle_elems = swizzle // element_bytewidth
+    if minor_dim % swizzle_elems == 0:
+      minor_tiling = swizzle_elems
+      break
+  else:
+    # No valid tile transform can be inferred.
+    raise ValueError(
+        f"{ref_ty.shape} is not a valid WGMMA shape"
+    )
+
+  return ir.ArrayAttr.get([
+      mgpu.TileTransformAttr.get((major_tiling, minor_tiling)),
+      mgpu.SwizzleTransformAttr.get(minor_tiling * element_bytewidth),
+  ])
+
+
+@partial(_add_transform_inference_rule, mgpu.WGMMAOp)
+def infer_wgmma_transforms(op: mgpu.WGMMAOp) -> OptionalTransforms:
+  b_transforms = infer_transforms_for_wgmma_ref(ir.MemRefType(op.b.type))
+  if ir.MemRefType.isinstance(op.a.type):
+    a_transforms = infer_transforms_for_wgmma_ref(
+        cast(ir.MemRefType, op.a.type)
+    )
+    return [a_transforms, b_transforms], []
+  return [b_transforms], []
+
+
+def _should_have_transforms(op: ir.OpView) -> bool:
+  """Returns 'True' if the operation should be assigned in/out transforms."""
+  return any(
+      map(
+          inference_utils.is_transformable_smem_memref,
+          itertools.chain(op.operands, op.results),
+      )
+  )
+
+
+def infer_transforms(module: ir.Module):
+  """Infers transforms for the given module.
+
+  Transforms are to smem refs the what layouts are to vectors.
+
+  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
+  annotate the same memref.
+  """
+
+  def inference_step(op: ir.Operation):
+    if not _should_have_transforms(op):
+      return
+    elif inference_rule := _transform_inference_rules.get(op.OPERATION_NAME, None):  # pytype: disable=attribute-error
+      pass
+    else:
+      raise NotImplementedError(f"Can not infer transforms for {op}")
+
+    maybe_transforms = inference_rule(op)
+    if maybe_transforms is None:
+      return
+
+    _set_transform_attributes(op, *maybe_transforms)
+
+  # It's enough to do a single backwards propagation (starting from vector
+  # users), and then a single forward propagation (to feed into the async loads
+  # and stores).
+  for op in module.body:
+    inference_utils.traverse_op(
+        op, inference_step, inference_utils.TraversalOrder.BACKWARDS
+    )
+  for op in module.body:
+    inference_utils.traverse_op(
+        op, inference_step, inference_utils.TraversalOrder.FORWARD
+    )


Test patch:
diff --git a/tests/mosaic/BUILD b/tests/mosaic/BUILD
index 4ec2dbf3b..71b2b7d80 100644
--- a/tests/mosaic/BUILD
+++ b/tests/mosaic/BUILD
@@ -68,6 +68,16 @@ jax_py_test(
     ] + py_deps("absl/testing"),
 )
 
+jax_py_test(
+    name = "gpu_transform_inference_test",
+    srcs = ["gpu_transform_inference_test.py"],
+    deps = [
+        "//jax",
+        "//jax:mosaic_gpu",
+        "//jax:test_util",
+    ] + py_deps("absl/testing"),
+)
+
 jax_multiplatform_test(
     name = "matmul_test",
     srcs = ["matmul_test.py"],
diff --git a/tests/mosaic/gpu_transform_inference_test.py b/tests/mosaic/gpu_transform_inference_test.py
new file mode 100644
index 000000000..2bae3fcd0
--- /dev/null
+++ b/tests/mosaic/gpu_transform_inference_test.py
@@ -0,0 +1,97 @@
+# Copyright 2025 The JAX Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Transform inference tests for the Mosaic GPU MLIR dialect."""
+
+# pylint: disable=g-complex-comprehension
+
+from absl.testing import parameterized
+import jax
+from jax import numpy as jnp
+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 func
+import jax.experimental.mosaic.gpu as mgpu
+from jax.experimental.mosaic.gpu import inference_utils
+import numpy as np
+
+
+config.parse_flags_with_absl()
+
+
+def _make_ir_context():
+  context = ir.Context()
+  context.append_dialect_registry(mlir_interpreter.upstream_dialects)
+  context.load_all_available_dialects()
+  mgpu.dialect.register_dialect(context)
+  return context
+
+
+class TransformInferenceTest(parameterized.TestCase):
+
+  def setUp(self):
+    if jax.version._version != jax.lib.__version__:
+      raise self.skipTest("Test requires matching jax and jaxlib versions")
+    super().setUp()
+    self.enter_context(_make_ir_context())
+    self.enter_context(ir.Location.unknown())
+    self.module = ir.Module.create()
+
+  @parameterized.parameters(
+      (swizzle, dtype)
+      for swizzle in mgpu.dialect.SwizzlingMode
+      for dtype in [jnp.bfloat16, jnp.float32]
+  )
+  def test_infer_transforms_for_wgmma_op(self, swizzle, dtype):
+    swizzle_elems = swizzle // np.dtype(dtype).itemsize
+    m = 64
+    # Note: `group_m` and `group_k` should be coprime with 2 for the test to be
+    # correct. Otherwise, we may infer larger swizzles than the test intends to
+    # check.
+    group_m, group_k = 3, 3
+    lhs_shape = (group_m * m, group_k * swizzle_elems)
+    rhs_shape = (group_k * swizzle_elems, group_k * swizzle_elems)
+    out_shape = (group_m * m, group_k * swizzle_elems)
+    wgmma_op = None
+
+    def body(accumulator, lhs, rhs):
+      nonlocal wgmma_op
+      wgmma_op = mgpu.dialect.WGMMAOp(accumulator, lhs, rhs)
+
+    with ir.InsertionPoint(self.module.body):
+      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+      elt_ty = mgpu.utils.dtype_to_ir_type(dtype)
+      lhs_ty = ir.MemRefType.get(lhs_shape, elt_ty, memory_space=smem)
+      rhs_ty = ir.MemRefType.get(rhs_shape, elt_ty, memory_space=smem)
+      acc_ty = ir.VectorType.get(out_shape, elt_ty)
+      func.FuncOp.from_py_func(acc_ty, lhs_ty, rhs_ty)(body)
+
+    mgpu.infer_transforms(self.module)
+
+    arg_transforms = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((8, swizzle_elems)),
+        mgpu.dialect.SwizzleTransformAttr.get(int(swizzle)),
+    ])
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(wgmma_op),
+        [arg_transforms, arg_transforms],
+    )
+    self.assertEmpty(inference_utils.out_transforms(wgmma_op))
+
+
+if __name__ == "__main__":
+  parameterized.absltest.main(testLoader=jtu.JaxTestLoader())