Repository: jax-ml/jax
Pull request #29006 (2025-05-26 09:51:49)

[Mosaic GPU] Resolve different tile transforms using the largest common divisor.
[Mosaic GPU] Resolve different tile transforms using the largest common divisor.

Before this change, transform inference would fail if different tilings needed to be used together (e.g. `(4 ,64)` and `(8, 32)`). After this change such tilings are resolved to use a tiling compatible with both, by taking the largest common divisor of each dimension. In the case above, the final tiling would be `(4, 32)`.

This change also adds two additional traversals in the transform inference, as this was necessary in a real-world example of tile resolution (taken from a ragged dot kernel).

The new test exercises both changes.

This CL also moves `is_known_divisible` from pallas to mgpu and extends it.



Patch:
diff --git a/jax/_src/pallas/mosaic_gpu/core.py b/jax/_src/pallas/mosaic_gpu/core.py
index 11f3f1eb4..7c27ae429 100644
--- a/jax/_src/pallas/mosaic_gpu/core.py
+++ b/jax/_src/pallas/mosaic_gpu/core.py
@@ -212,30 +212,6 @@ def kernel(
   return wrapper
 
 
-def _is_known_divisible(value, divisor, fuel=10) -> bool:
-  """Returns True if the value is statically known to be divisible by the divisor."""
-  if divisor == 1:
-    return True
-  if fuel < 0:
-    return False
-  if not isinstance(value.owner, ir.Operation):
-    return False
-  def_op = value.owner.opview
-  match def_op:
-    case arith_dialect.IndexCastOp():
-      return _is_known_divisible(value.owner.operands[0], divisor, fuel - 1)
-    case arith_dialect.ConstantOp():
-      return ir.IntegerAttr(def_op.value).value % divisor == 0
-    case arith_dialect.MulIOp():
-      return (_is_known_divisible(value.owner.operands[0], divisor, fuel // 2) or
-              _is_known_divisible(value.owner.operands[1], divisor, (fuel + 1)// 2))
-    case arith_dialect.SelectOp():
-      return (_is_known_divisible(value.owner.operands[1], divisor, fuel // 2) and
-              _is_known_divisible(value.owner.operands[2], divisor, (fuel + 1)// 2))
-
-  return False
-
-
 @dataclasses.dataclass(frozen=True)
 class GPUMemoryRef(pallas_core.MemoryRef):
   transforms: Sequence[MemoryRefTransform] = ()
@@ -503,7 +479,7 @@ class UntileRef(state_types.Transform):
               f" tiling ({tile})"
           )
         if isinstance(idx.base, ir.Value):
-          if not _is_known_divisible(idx.base, tile):
+          if not mgpu_utils.is_known_divisible(idx.base, tile):
             raise ValueError(
                 "Dynamic slice base index (which is a dynamic value) cannot be"
                 f" statically proven to be divisible by the tiling ({tile})"
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index dbee817df..7094fc735 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -86,6 +86,7 @@ from .utils import (
     debug_print as debug_print,
     ds as ds,
     fori as fori,
+    is_known_divisible as is_known_divisible,
     memref_fold as memref_fold,
     memref_slice as memref_slice,
     memref_reshape as memref_reshape,
diff --git a/jax/experimental/mosaic/gpu/transform_inference.py b/jax/experimental/mosaic/gpu/transform_inference.py
index c8a4afd61..a39ba2899 100644
--- a/jax/experimental/mosaic/gpu/transform_inference.py
+++ b/jax/experimental/mosaic/gpu/transform_inference.py
@@ -20,6 +20,7 @@ annotated with layouts (see `layout_inference.py` for the relevant pass).
 
 from collections.abc import Callable
 from functools import partial
+import math
 from typing import cast
 
 from jax._src import lib as jaxlib
@@ -82,12 +83,28 @@ def _resolve_transforms(
   if other_transforms is None:
     return transforms
 
-  if transforms != other_transforms:
+  if len(transforms) != len(other_transforms):
     raise NotImplementedError(
         f"Conflicting transforms {transforms} != {other_transforms}."
     )
 
-  return transforms
+  new_transforms = []
+  for a, b in zip(transforms, other_transforms, strict=True):
+    if a == b:
+      new_transforms.append(a)
+    elif mgpu.TileTransformAttr.isinstance(a) and mgpu.TileTransformAttr.isinstance(b):
+      a = mgpu.TileTransformAttr(a)
+      b = mgpu.TileTransformAttr(b)
+      if len(a.tiling) != len(b.tiling):
+        raise ValueError(f"Conflicting tile transforms {a} != {b}.")
+      new_tiling = []
+      for tile_a, tile_b in zip(a.tiling, b.tiling):
+        new_tiling.append(math.gcd(tile_a, tile_b))
+      new_transforms.append(mgpu.TileTransformAttr.get(new_tiling))
+    else:
+      raise NotImplementedError(f"Unsupported transforms {a} and {b}")
+
+  return ir.ArrayAttr.get(new_transforms)
 
 
 def _transforms_from_uses(op: ir.OpView) -> ir.Attribute | None:
@@ -280,7 +297,7 @@ def _infer_memref_subview_transforms(
   #  - We only propagate transforms if they consist of a single tile transform
   #    and a single swizzle transform.
   # TODO(bchetioui): implement more complex propagation rules.
-  tile_transform, _ = _get_tile_and_swizzle_transforms(transforms)
+  tile_transform, swizzle_transform = _get_tile_and_swizzle_transforms(transforms)
 
   # Check swizzle transform propagation.
   strides, _ = ir.MemRefType.get_strides_and_offset(op.source.type)
@@ -292,9 +309,41 @@ def _infer_memref_subview_transforms(
     )
 
   # Check tile transform propagation.
-  # TODO(dasenov): implement more precise checks.
+  old_tiling = mgpu.TileTransformAttr(tile_transform).tiling
+  num_tiled_axes = len(old_tiling)
+  last_n_dims = op.source.type.shape[-num_tiled_axes:]
+  last_n_sizes = list(op.static_sizes)[-num_tiled_axes:]
+  last_n_offsets = list(op.static_offsets)[-num_tiled_axes:]
 
-  return [transforms], [transforms]
+  if any(ir.ShapedType.is_dynamic_size(x) for x in last_n_sizes):
+    raise NotImplementedError(
+        "Subview transforms with dynamic sizes are not supported."
+    )
+
+  dynamic_index = 0
+  for i in range(len(last_n_offsets)):
+    if ir.ShapedType.is_dynamic_size(last_n_offsets[i]):
+      if utils.is_known_divisible(
+          op.offsets[dynamic_index], last_n_sizes[i]
+      ):
+        last_n_offsets[i] = last_n_sizes[i]
+      else:
+        # This will force a tiling of 1 along this axis. This is a safe choice
+        # (since we couldn't infer a better one) but might not be optimal.
+        last_n_offsets[i] = 1
+      dynamic_index += 1
+
+  new_tiling = [
+      math.gcd(*xs)
+      for xs in zip(
+          last_n_sizes, last_n_dims, last_n_offsets, old_tiling, strict=True
+      )
+  ]
+
+  new_transforms = ir.ArrayAttr.get(
+      [mgpu.TileTransformAttr.get(new_tiling), swizzle_transform]
+  )
+  return [new_transforms], [new_transforms]
 
 
 @partial(_add_transform_inference_rule, memref.TransposeOp)
@@ -381,17 +430,20 @@ def infer_transforms(module: ir.Module):
 
     _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
-    )
+  # We alternate a few backwards propagation (starting from vector users), and
+  # forward propagation (to feed into the async loads and stores) passes in
+  # order to enable more complex inference situations.
+  #
+  # TODO(bchetioui): Replace this with a more generic inference.
+  inference_passes = [
+      inference_utils.TraversalOrder.BACKWARDS,
+      inference_utils.TraversalOrder.FORWARD,
+      inference_utils.TraversalOrder.BACKWARDS,
+      inference_utils.TraversalOrder.FORWARD,
+  ]
+  for traversal_order in inference_passes:
+    for op in module.body:
+      inference_utils.traverse_op(op, inference_step, traversal_order)
 
   # All ops that should have transforms but have no transforms inferred so far
   # are assigned an empty sets of transforms. E.g., this happens in kernels with
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index 9ea964c98..678aad0c9 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -1498,3 +1498,43 @@ def vector_concat(vectors: Sequence[ir.Value]) -> ir.Value:
       result = vector.insertelement(elem, result, position=c(offset + i, index))
     offset += vty.shape[0]
   return result
+
+
+def is_known_divisible(value, divisor, max_depth=10) -> bool:
+  """Returns True if the value is statically known to be divisible by the divisor."""
+  if divisor == 1:
+    return True
+  if max_depth < 0 or not isinstance(value.owner, ir.Operation):
+    return False
+
+  new_depth = max_depth - 1
+  def_op = value.owner.opview
+
+  match def_op:
+    case arith.IndexCastOp():
+      return is_known_divisible(value.owner.operands[0], divisor, max_depth - 1)
+    case arith.ConstantOp():
+      return ir.IntegerAttr(def_op.value).value % divisor == 0
+    case arith.MulIOp():
+      # Only cover the case where one operand is divisible. It's still possible
+      # that the final product is divisible, but we don't check that here.
+      return (is_known_divisible(value.owner.operands[0], divisor, new_depth) or
+              is_known_divisible(value.owner.operands[1], divisor, new_depth))
+    case arith.SelectOp():
+      return (is_known_divisible(value.owner.operands[1], divisor, new_depth) and
+              is_known_divisible(value.owner.operands[2], divisor, new_depth))
+    case arith.MaxSIOp() | arith.MinSIOp() | arith.MaxUIOp() | arith.MinUIOp():
+      return (is_known_divisible(value.owner.operands[0], divisor, new_depth) and
+              is_known_divisible(value.owner.operands[1], divisor, new_depth))
+    case arith.AddIOp() | arith.SubIOp():
+      # Only cover the common case where both operads are divisible.
+      return (is_known_divisible(value.owner.operands[0], divisor, new_depth) and
+              is_known_divisible(value.owner.operands[1], divisor, new_depth))
+    case arith.AndIOp():
+      # Only cover the specific case where the divisor is a power of two.
+      return divisor.bit_count() == 1 and (
+          is_known_divisible(value.owner.operands[0], divisor, new_depth)
+          or is_known_divisible(value.owner.operands[1], divisor, new_depth)
+      )
+
+  return False


Test patch:
diff --git a/tests/mosaic/gpu_transform_inference_test.py b/tests/mosaic/gpu_transform_inference_test.py
index 420d735fe..b4a41a141 100644
--- a/tests/mosaic/gpu_transform_inference_test.py
+++ b/tests/mosaic/gpu_transform_inference_test.py
@@ -19,6 +19,7 @@
 from absl.testing import parameterized
 import jax
 from jax import numpy as jnp
+from jax._src import lib as jaxlib
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.interpreters import mlir as mlir_interpreter
@@ -543,6 +544,127 @@ class TransformInferenceTest(parameterized.TestCase):
     with self.assertRaises(NotImplementedError):
       mgpu.infer_transforms(self.module)
 
+  @parameterized.parameters([False, True])
+  def test_infer_transforms_for_sibling_subviews_and_distant_op(
+      self, even_offsets
+  ):
+    # This test uses the following op tree extracted from this ragged dot
+    # kernel:
+    # https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/gpu/ragged_dot_mgpu.py
+    #
+    #   subview_op0   (slice = 64, 64)
+    #   - subview_op1 (slice = 2, 64)
+    #   - subview_op2 (slice = 4, 64, either at an even or odd offset)
+    #   - subview_op3 (slice = 8, 64)
+    #   - user_op0    (in_transforms = [tile(64, 64), swizzle(32)])
+    #
+    # First the in_transforms of user_op0 have to be propagated up to
+    # subview_op0. Then they have to be propagated down and resolved. Finally
+    # all subview ops need to have the same transforms.
+
+    # TODO(dasenov): Remove this after the minimal jaxlib version is 0.6.2.
+    if jaxlib.version < (0, 6, 2):
+      self.skipTest("Test requires jaxlib version >= 0.6.2")
+
+    subview_op0, subview_op1, subview_op2, subview_op3 = None, None, None, None
+    user_op0 = None
+
+    source_shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+    smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
+    source_ref_ty = ir.MemRefType.get(source_shape, elt_ty, memory_space=smem)
+
+    slice1_shape = (2, 64)
+    slice2_shape = (4, 64)
+    slice3_shape = (8, 64)
+
+    slice0_ref_ty = ir.MemRefType.get(source_shape, elt_ty, memory_space=smem)
+    slice1_ref_ty = ir.MemRefType.get(slice1_shape, elt_ty, memory_space=smem)
+    slice2_ref_ty = ir.MemRefType.get(slice2_shape, elt_ty, memory_space=smem)
+    slice3_ref_ty = ir.MemRefType.get(slice3_shape, elt_ty, memory_space=smem)
+
+    def body(source_ref):
+      nonlocal subview_op0, subview_op1, subview_op2, subview_op3, user_op0
+
+      subview_op0 = memref.SubViewOp(
+          slice0_ref_ty,
+          source_ref,
+          [],  # dynamic offsets
+          [],  # dynamic sizes
+          [],  # dynamic strides
+          static_offsets=[0, 0],
+          static_sizes=source_shape,
+          static_strides=[1, 1],
+      )
+
+      transforms_0 = ir.ArrayAttr.get([
+          mgpu.dialect.TileTransformAttr.get((64, 64)),
+          mgpu.dialect.SwizzleTransformAttr.get(32),
+      ])
+      user_op0 = mgpu.dialect.WithTransformsOp(subview_op0.result, transforms_0)
+
+      subview_op1 = memref.SubViewOp(
+          slice1_ref_ty,
+          subview_op0,
+          [],  # dynamic offsets
+          [],  # dynamic sizes
+          [],  # dynamic strides
+          static_offsets=[0, 0],
+          static_sizes=slice1_shape,
+          static_strides=[1, 1],
+      )
+
+      subview_op2 = memref.SubViewOp(
+          slice2_ref_ty,
+          subview_op0,
+          [],  # dynamic offsets
+          [],  # dynamic sizes
+          [],  # dynamic strides
+          static_offsets=[16 if even_offsets else 15, 0],
+          static_sizes=slice2_shape,
+          static_strides=[1, 1],
+      )
+
+      # The following ops are just to test the dynamic offsets support.
+      c = lambda x: arith.constant(ir.IntegerType.get_signless(32), x)
+      c64 = c(64)
+      c32 = c(32)
+      c16 = c(16)
+      subi = arith.subi(c64, c32)
+      maxsi = arith.maxsi(c16, subi)
+      addi = arith.addi(maxsi, subi)
+      andi = arith.andi(addi, maxsi)
+      idx = arith.index_cast(ir.IndexType.get(), andi)
+      subview_op3 = memref.SubViewOp(
+          slice3_ref_ty,
+          subview_op0,
+          [idx],  # dynamic offsets
+          [],  # dynamic sizes
+          [],  # dynamic strides
+          static_offsets=[ir.ShapedType.get_dynamic_size(), 0],
+          static_sizes=slice3_shape,
+          static_strides=[1, 1],
+      )
+
+    with ir.InsertionPoint(self.module.body):
+      func.FuncOp.from_py_func(source_ref_ty)(body)
+
+    mgpu.infer_transforms(self.module)
+
+    want = ir.ArrayAttr.get([
+          mgpu.dialect.TileTransformAttr.get((2 if even_offsets else 1, 64)),
+          mgpu.dialect.SwizzleTransformAttr.get(32),
+      ])
+
+    self.assertSequenceEqual(inference_utils.in_transforms(subview_op0), [want])
+    self.assertSequenceEqual(inference_utils.out_transforms(subview_op0), [want])
+    self.assertSequenceEqual(inference_utils.in_transforms(subview_op1), [want])
+    self.assertSequenceEqual(inference_utils.out_transforms(subview_op1), [want])
+    self.assertSequenceEqual(inference_utils.in_transforms(subview_op2), [want])
+    self.assertSequenceEqual(inference_utils.out_transforms(subview_op2), [want])
+    self.assertSequenceEqual(inference_utils.in_transforms(subview_op3), [want])
+    self.assertSequenceEqual(inference_utils.out_transforms(subview_op3), [want])
+
 
 if __name__ == "__main__":
   parameterized.absltest.main(testLoader=jtu.JaxTestLoader())