[Mosaic GPU] Fix a bug with transform inference of subview in the presence of dynamic offsets.

The old code incorrectly started counting the dynamic offsets from 0 while processing the tiled dimensions. This only works if there are no dynamic offsets in the non-tiled dimensions. Otherwise, the logic will index into the wrong offset.

PiperOrigin-RevId: 802047327

Author: dimitar-asenov

jax/experimental/mosaic/gpu/transform_inference.py

@@ -334,21 +334,26 @@ def _infer_memref_subview_transforms(
         "unchanged."
     )
 
+  is_dynamic = lambda x: ir.ShapedType.is_dynamic_size(x)
+
   # Check tile transform propagation.
   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:]
 
-  if any(ir.ShapedType.is_dynamic_size(x) for x in last_n_sizes):
+  if any(is_dynamic(x) for x in last_n_sizes):
     raise NotImplementedError(
         "Subview transforms with dynamic sizes are not supported."
     )
 
-  dynamic_index = 0
+  num_non_tiled_axes = len(op.source.type.shape) - num_tiled_axes
+  non_tiled_offsets = list(op.static_offsets)[:num_non_tiled_axes]
+  dynamic_index = sum(1 for x in non_tiled_offsets if is_dynamic(x))
+
   for i in range(len(last_n_offsets)):
-    if ir.ShapedType.is_dynamic_size(last_n_offsets[i]):
+    if is_dynamic(last_n_offsets[i]):
       if utils.is_known_divisible(
           op.offsets[dynamic_index], last_n_sizes[i]
       ):

tests/mosaic/gpu_transform_inference_test.py

@@ -657,6 +657,62 @@ def test_custom_primitive_op_retains_transforms(self):
     mgpu.infer_transforms(self.module)
     self.assertSequenceEqual(inference_utils.in_transforms(op), [transforms])
 
+  @parameterized.parameters([False, True])
+  def test_infer_transforms_for_subview_handles_dynamic_offsets(
+      self, annotate_input
+  ):
+    subview_op = user_op = None
+    shape = (32, 32, 32)
+    elt_ty = ir.BF16Type.get()
+
+    in_ref_ty = ir.MemRefType.get(shape, elt_ty, memory_space=mgpu.utils.smem())
+    out_ref_ty = ir.MemRefType.get((16, 16, 32), elt_ty, memory_space=mgpu.utils.smem())
+
+    def body(in_ref):
+      nonlocal subview_op, user_op
+      c = lambda x: arith.constant(ir.IntegerType.get_signless(32), x)
+      subview_op = memref.SubViewOp(
+          out_ref_ty,
+          in_ref,
+          [c(16), c(4)],
+          [],
+          [],
+          static_offsets=[
+              ir.ShapedType.get_dynamic_size(),
+              ir.ShapedType.get_dynamic_size(),
+              0,
+          ],
+          static_sizes=[16, 16, 32],
+          static_strides=[1, 1, 1],
+      )
+      user_op = memref.CastOp(out_ref_ty, subview_op.result)
+
+    with ir.InsertionPoint(self.module.body):
+      f = func.FuncOp.from_py_func(in_ref_ty)(body).func_op
+
+    transforms = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((16, 16)),
+        mgpu.dialect.SwizzleTransformAttr.get(32),
+    ])
+
+    if annotate_input:
+      f.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+    else:
+      user_op.attributes["in_transforms"] = ir.ArrayAttr.get([transforms])
+
+    mgpu.infer_transforms(self.module)
+
+    expected_transforms = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((1, 16)),
+        mgpu.dialect.SwizzleTransformAttr.get(32),
+    ])
+
+    self.assertSequenceEqual(
+        inference_utils.in_transforms(subview_op), [expected_transforms]
+    )
+    self.assertSequenceEqual(
+        inference_utils.out_transforms(subview_op), [expected_transforms]
+    )
 
 if __name__ == "__main__":
   parameterized.absltest.main(testLoader=jtu.JaxTestLoader())