[Mosaic GPU][NFC] Specify the memory space for two existing tests.

This will become relevant later once we enable SMEM inference by default.

PiperOrigin-RevId: 813743079

Author: dimitar-asenov

jax/experimental/mosaic/gpu/layout_inference.py

@@ -1125,6 +1125,17 @@ def _async_store_tmem_equation_system(
   )
 
 
+@_add_equation_system_derivation_rule(mgpu.SliceSMEMOp)
+def _slice_smem_equation_system(
+    ctx: DerivationContext,
+    op: mgpu.SliceSMEMOp,
+) -> tuple[eqns.EquationSystem, OperandOrResultsForVariable, list[Hint]]:
+  del ctx
+  res = OperandOrResult(op, VariableType.RESULT, 0)
+  res_var = eqns.Variable(res)
+  return (eqns.EquationSystem(), {res_var: [res]}, [])
+
+
 # `memref.load` and `memref.store` are used to load barrier phases which are
 # scalars---the rule needn't do anything interesting, but we need to have it.
 @_add_equation_system_derivation_rule(memref.LoadOp)
@@ -1148,6 +1159,43 @@ def _memref_load_store_op_equation_system(
   return eqns.EquationSystem(assignments=assignments), {var: [ref]}, []
 
 
+@_add_equation_system_derivation_rule(mgpu.WithTransformsOp)
+def _with_transforms_equation_system(
+    ctx: DerivationContext,
+    op: mgpu.WithTransformsOp,
+) -> tuple[eqns.EquationSystem, OperandOrResultsForVariable, list[Hint]]:
+  source = OperandOrResult(op, VariableType.OPERAND, 0)
+  dest = OperandOrResult(op, VariableType.RESULT, 0)
+  var = ctx.producer_ref(source)
+
+  transforms = [layouts_lib.from_transform_attr(x) for x in op.transforms]
+  match transforms:
+    case []:
+      tile_transform = None
+      swizzle = None
+    case [lc.TileTransform() as t]:
+      tile_transform = t
+      swizzle = None
+    case [lc.TileTransform() as t, mgpu.SwizzlingMode() as s]:
+      tile_transform = t
+      swizzle = s
+    case _:
+      raise NotImplementedError(f"Unsupported transforms {transforms}")
+
+  if swizzle is not None:
+    computed_swizzle = _compute_swizzle(op.ref.type, tile_transform)
+    if computed_swizzle != swizzle:
+      raise NotImplementedError(
+          f"Cannot honor caller-provided swizzle {swizzle} that is different "
+          f"from the computed swizle {computed_swizzle} on op {op}."
+      )
+
+  assignments: dict[eqns.Variable, eqns.Constant] = {
+      var: eqns.SMEMTiling(tile_transform)
+  }
+  return eqns.EquationSystem(assignments=assignments), {var: [source, dest]}, []
+
+
 def _ensure_all_layouts_are_set(
     op: ir.OpView, enable_smem_inference: bool
 ) -> None:

tests/mosaic/gpu_layout_inference_test.py

@@ -413,7 +413,9 @@ def test_infer_layout_from_body_op_to_yield_op_to_for_op(self):
     shape = (64, 64)
     with ir.InsertionPoint(self.module.body):
       c_ty = ir.VectorType.get(shape, ir.BF16Type.get())
-      ab_type = ir.MemRefType.get(shape, ir.BF16Type.get())
+      ab_type = ir.MemRefType.get(
+          shape, ir.BF16Type.get(), memory_space=mgpu.utils.smem()
+      )
       i32 = ir.IntegerType.get_signless(32)
       lower_bound, upper_bound, step, a, b, c = undefs(
           i32, i32, i32, ab_type, ab_type, c_ty
@@ -794,7 +796,7 @@ def test_infer_wgmma_layout_correctly(self, lhs_memory_space):
 
     with ir.InsertionPoint(self.module.body):
       vec_ty = ir.VectorType.get(shape, f32)
-      ref_ty = ir.MemRefType.get(shape, f32)
+      ref_ty = ir.MemRefType.get(shape, f32, memory_space=mgpu.utils.smem())
       lhs_ty = ref_ty if lhs_memory_space == "shared" else vec_ty
       acc, lhs, rhs = undefs(vec_ty, lhs_ty, ref_ty)
       wgmma_op = mgpu.dialect.WGMMAOp(acc, lhs, rhs)
@@ -1253,6 +1255,62 @@ def test_memref_load_store_op_transforms_are_empty(self):
       self.assertEqual(inference_utils.in_transforms(load_op), want)
       self.assertEqual(inference_utils.in_transforms(store_op), want)
 
+  def test_slice_smem_gets_empty_by_default(self):
+    with ir.InsertionPoint(self.module.body):
+      shape = (64, 64)
+      elt_ty = ir.BF16Type.get()
+      i32 = ir.IntegerType.get_signless(32)
+      [offset] = undefs(i32)
+      ref_ty = ir.MemRefType.get(shape, elt_ty, memory_space=mgpu.utils.smem())
+      slice_smem_op = mgpu.dialect.SliceSMEMOp(ref_ty, offset)
+
+      transforms = ir.ArrayAttr.get([])
+      mgpu.infer_layout(self.module, enable_smem_inference=True)
+      self.assertSequenceEqual(
+          inference_utils.out_transforms(slice_smem_op), [transforms]
+      )
+
+  def test_infer_transforms_preserves_with_transforms_requirements(self):
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    with ir.InsertionPoint(self.module.body):
+      ref_ty = ir.MemRefType.get(shape, elt_ty, memory_space=mgpu.utils.smem())
+      [ref] = undefs(ref_ty)
+
+      transforms = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((8, 64)),
+        mgpu.dialect.SwizzleTransformAttr.get(128),
+      ])
+      mgpu.dialect.with_transforms(ref, transforms)
+
+    mgpu.infer_layout(self.module, enable_smem_inference=True)
+    self.assertSequenceEqual(
+        inference_utils.out_transforms(ref.owner), [transforms]
+    )
+
+  def test_infer_transforms_fails_on_conflicting_with_transforms_requirements(self):
+    shape = (64, 64)
+    elt_ty = ir.BF16Type.get()
+
+    with ir.InsertionPoint(self.module.body):
+      ref_ty = ir.MemRefType.get(shape, elt_ty, memory_space=mgpu.utils.smem())
+      [ref] = undefs(ref_ty)
+
+      transforms1 = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((8, 64)),
+        mgpu.dialect.SwizzleTransformAttr.get(128),
+      ])
+      transforms2 = ir.ArrayAttr.get([
+        mgpu.dialect.TileTransformAttr.get((16, 64)),
+        mgpu.dialect.SwizzleTransformAttr.get(128),
+      ])
+      mgpu.dialect.with_transforms(ref, transforms1)
+      mgpu.dialect.with_transforms(ref, transforms2)
+
+    with self.assertRaisesRegex(ValueError, "Failed to infer"):
+      mgpu.infer_layout(self.module, enable_smem_inference=True)
+
 
 if __name__ == "__main__":
   parameterized.absltest.main(testLoader=jtu.JaxTestLoader())