[Mosaic GPU][NFC] Add utils to create and check the SMEM memory space attribute.

PiperOrigin-RevId: 781447752

Author: dimitar-asenov

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -494,12 +494,13 @@ def scratch_view(
       runtime scratch buffer.
     """
     smem_base = None
-    smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
     i8 = ir.IntegerType.get_signless(8)
     i32 = ir.IntegerType.get_signless(32)
     if self.lowering_semantics == mgpu.LoweringSemantics.Lane:
       smem_base = gpu_dialect.dynamic_shared_memory(
-          ir.MemRefType.get((mgpu_utils.DYNAMIC,), i8, memory_space=smem)
+          ir.MemRefType.get(
+              (mgpu_utils.DYNAMIC,), i8, memory_space=mgpu_utils.smem()
+          )
       )
     views = []
     off = initial_used_bytes = self.smem_used_bytes
@@ -508,7 +509,7 @@ def scratch_view(
       scratch_ty = ir.MemRefType.get(
           s.shape,
           mgpu_utils.dtype_to_ir_type(s.dtype),
-          memory_space=smem,
+          memory_space=mgpu_utils.smem(),
       )
       # The below code emission relies on the assumption that the first scratch
       # operand provided by Mosaic GPU always begins at the beginning of
@@ -1271,7 +1272,7 @@ def _handle_dtype_bitcast(
         "Data type bitcast is only supported from i8 to other types."
     )
   ref_ty = ir.MemRefType(ref.type)
-  if ref_ty.memory_space != ir.Attribute.parse("#gpu.address_space<workgroup>"):
+  if not mgpu_utils.is_smem_ref(ref_ty):
     raise ValueError(f"Only workgroup memory is supported but got {ref}.")
   if len(ref_ty.shape) != 1:
     raise NotImplementedError(

jax/experimental/mosaic/gpu/core.py

@@ -350,7 +350,6 @@ def _construct_smem_reftree(
   index = ir.IndexType.get()
   i32 = ir.IntegerType.get_signless(32)
   i64 = ir.IntegerType.get_signless(64)
-  smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
   flat_ref_tys, smem_buffer_tree = jax.tree.flatten(
       smem_buffers, is_leaf=lambda x: isinstance(x, Union)
   )
@@ -364,7 +363,7 @@ def barrier_memref(num_barriers: int) -> ir.Value:
           ir.Type.parse("!mosaic_gpu.barrier")
           if lowering_semantics == LoweringSemantics.Warpgroup
           else i64,
-          memory_space=smem,
+          memory_space=utils.smem(),
       )
       barrier_memref = _slice_smem(
             barrier_ty,
@@ -411,7 +410,7 @@ def ref(member_thunks=member_thunks):
         )
       case TMEM(shape, dtype, layout=layout, collective=collective, packing=packing):
         addr_ref = _slice_smem(
-            ir.MemRefType.get([], i32, memory_space=smem),
+            ir.MemRefType.get([], i32, memory_space=utils.smem()),
             dynamic_smem,
             c(dynamic_smem_offset, index),
             lowering_semantics,
@@ -433,7 +432,7 @@ def ref(addr_ref=addr_ref, shape=shape, dtype=dtype, layout=layout):
       case _:
         mlir_dtype = utils.dtype_to_ir_type(ref_ty.dtype)
         tile_smem = _slice_smem(
-            ir.MemRefType.get(ref_ty.shape, mlir_dtype, memory_space=smem),
+            ir.MemRefType.get(ref_ty.shape, mlir_dtype, memory_space=utils.smem()),
             dynamic_smem,
             c(dynamic_smem_offset, index),
             lowering_semantics,
@@ -531,18 +530,17 @@ def _launch(
       token.type, [token], *grid_vals, *block_vals,
       dynamicSharedMemorySize=c(smem_bytes, i32), **cluster_kwargs)
   launch_op.body.blocks.append(*([index] * (12 + 2 * len(cluster_kwargs))))  # Append an empty block
-  smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
   with ir.InsertionPoint(launch_op.body.blocks[0]):
     dynamic_smem = gpu.dynamic_shared_memory(
-        ir.MemRefType.get((utils.DYNAMIC,), i8, memory_space=smem)
+        ir.MemRefType.get((utils.DYNAMIC,), i8, memory_space=utils.smem())
     )
 
     if profiler_spec:
       prof_smem = _slice_smem(
           ir.MemRefType.get(
               (profiler_spec.smem_i32_elements(block=block),),
               i32,
-              memory_space=smem,
+              memory_space=utils.smem(),
           ),
           dynamic_smem,
           c(profiler_start, index),

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -514,9 +514,8 @@ def _vector_reduction_op_lowering_rule(
   a = _fragmented_array_from_ir(op.vector, layout, is_signed)
   match str(op.kind):
     case "#vector.kind<add>":
-      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
       scratch = _slice_smem(
-          ir.MemRefType.get([4], element_type, memory_space=smem),
+          ir.MemRefType.get([4], element_type, memory_space=utils.smem()),
           arith.constant(None, op.attributes["offset"]),
       )
       result = a.reduce("add", range(len(a.shape)), scratch)
@@ -675,7 +674,7 @@ def _transformed_smem_ref_type(
   if not transforms and not transposed:
     return ref_ty
 
-  if ref_ty.memory_space != ir.Attribute.parse("#gpu.address_space<workgroup>"):
+  if not utils.is_smem_ref(ref_ty):
     raise ValueError(f"Only workgroup memory is supported but got {ref_ty}.")
 
   shape = ref_ty.shape
@@ -1130,17 +1129,16 @@ def _mgpu_slice_smem_op_lowering_rule(
 
 def _slice_smem(result: ir.Type, offset: ir.Value):
   i8 = ir.IntegerType.get_signless(8)
-  smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
   smem_base = gpu.dynamic_shared_memory(
-      ir.MemRefType.get((utils.DYNAMIC,), i8, memory_space=smem)
+      ir.MemRefType.get((utils.DYNAMIC,), i8, memory_space=utils.smem())
   )
   offset = arith.index_cast(ir.IndexType.get(), offset)
   lowered_result_type = result
   if ir.MemRefType.isinstance(result):
     memref_ty = ir.MemRefType(result)
     if memref_ty.element_type == ir.Type.parse("!mosaic_gpu.barrier"):
       lowered_result_type = ir.MemRefType.get(
-          memref_ty.shape, _lowered_barrier_type(), memory_space=smem
+          memref_ty.shape, _lowered_barrier_type(), memory_space=utils.smem()
       )
   view = memref.view(lowered_result_type, smem_base, offset, [])
   if result == lowered_result_type:

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -2665,7 +2665,7 @@ def transfer_tiled2(
 
     if ref_ty.memory_space is None:
       llvm_memory_space = None
-    elif ref_ty.memory_space == ir.Attribute.parse("#gpu.address_space<workgroup>"):
+    elif utils.is_smem_ref(ref_ty):
       llvm_memory_space = 3
     else:
       raise ValueError(f"Unsupported memory space: {ref_ty.memory_space}")

jax/experimental/mosaic/gpu/inference_utils.py

@@ -22,6 +22,8 @@
 
 from jax._src.lib.mlir import ir
 
+from . import utils
+
 MlirOperation = Union[ir.Operation, ir.OpView]
 
 def in_layouts(op: MlirOperation) -> Sequence[ir.Attribute]:
@@ -147,13 +149,11 @@ def should_have_transforms(op: ir.OpView) -> bool:
 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
+      and utils.is_smem_ref(v)
   )
 
 

jax/experimental/mosaic/gpu/launch_context.py

@@ -594,7 +594,6 @@ def async_copy(
     index = ir.IndexType.get()
     i16 = ir.IntegerType.get_signless(16)
     i32 = ir.IntegerType.get_signless(32)
-    smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
     src_ref_ty = ir.MemRefType(src_ref.type)
     dst_ref_ty = ir.MemRefType(dst_ref.type)
     element_type = src_ref_ty.element_type
@@ -609,13 +608,13 @@ def async_copy(
     if not isinstance(gmem_transform, tuple):
       gmem_transform = (gmem_transform,)
 
-    if src_ref_ty.memory_space is None and dst_ref_ty.memory_space == smem:
+    if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty):
       gmem_ref, smem_ref = src_ref, dst_ref
       if barrier is None:
         raise ValueError("Barriers are required for GMEM -> SMEM copies")
       if arrive is None:
         arrive = True  # Arrive by default
-    elif src_ref_ty.memory_space == smem and dst_ref_ty.memory_space is None:
+    elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None:
       gmem_ref, smem_ref = dst_ref, src_ref
       if barrier is not None:
         raise ValueError("Barriers are unsupported for SMEM -> GMEM copies")

jax/experimental/mosaic/gpu/tcgen05.py

@@ -613,7 +613,7 @@ def tmem_alloc(tmem_addr: ir.Value, ncols: int, collective: bool = False, exact:
     ref_ty = ir.MemRefType(tmem_addr.type)
     if ref_ty.element_type != ir.IntegerType.get_signless(32):
       raise ValueError(f"tmem_addr must be an i32 memref, got: {ref_ty}")
-    if ref_ty.memory_space != ir.Attribute.parse("#gpu.address_space<workgroup>"):
+    if not utils.is_smem_ref(ref_ty):
       raise ValueError(f"tmem_addr must be in shared memory, got: {ref_ty}")
     if math.prod(ref_ty.shape) != 1:
       raise ValueError(f"tmem_addr must contain a single element, got: {ref_ty}")
@@ -870,9 +870,8 @@ def from_alloc(
     if not ir.MemRefType.isinstance(tmem_addr_ref.type):
       raise ValueError(f"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}")
     addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)
-    smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
-    if addr_ref_ty.memory_space != smem:
-      raise ValueError(f"tmem_addr_ref must be in workgroup memory, got: {addr_ref_ty}")
+    if not utils.is_smem_ref(addr_ref_ty):
+      raise ValueError(f"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}")
     if addr_ref_ty.element_type != i32:
       raise ValueError(f"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}")
     if math.prod(addr_ref_ty.shape) != 1:

jax/experimental/mosaic/gpu/utils.py

@@ -1548,3 +1548,20 @@ def is_known_divisible(value, divisor, max_depth=10) -> bool:
       )
 
   return False
+
+
+def smem() -> ir.Attribute:
+  """Returns the attribute for the SMEM memory space."""
+  return ir.Attribute.parse("#gpu.address_space<workgroup>")
+
+
+def is_smem_ref(ref: ir.Value | ir.Type) -> bool:
+  """Returns true if the input mem ref or memref type points to SMEM.
+  If the input is not at all of a memref type, raises a ValueError.
+  """
+  if isinstance(ref, ir.Value):
+    ref = ref.type
+  if not ir.MemRefType.isinstance(ref):
+    raise ValueError(f"Expected a memref type but got {ref}")
+  ref = ir.MemRefType(ref)
+  return ref.memory_space is not None and ref.memory_space == smem()

tests/mosaic/gpu_dialect_test.py

@@ -967,8 +967,7 @@ def test_lowering_slice_smem_op(self):
     def body():
       nonlocal offset
       i32 = ir.IntegerType.get_signless(32)
-      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
-      memref_ty = ir.MemRefType.get((4, 32), i32, memory_space=smem)
+      memref_ty = ir.MemRefType.get((4, 32), i32, memory_space=mgpu_utils.smem())
       offset = arith.constant(i32, shift)
       op = mgpu.dialect.SliceSMEMOp(memref_ty, offset)
       op.attributes["out_transforms"] = ir.ArrayAttr.get([ir.ArrayAttr.get([])])
@@ -1051,8 +1050,7 @@ def body(vec1, vec2, ref):
       )
 
     with ir.InsertionPoint(self.module.body):
-      smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
-      ref_ty = ir.MemRefType.get((4, 32), ir.BF16Type.get(), memory_space=smem)
+      ref_ty = ir.MemRefType.get((4, 32), ir.BF16Type.get(), memory_space=mgpu_utils.smem())
       func.FuncOp.from_py_func(vec_ty, vec_ty, ref_ty)(body)
 
     if omit_in_layouts:
@@ -1071,7 +1069,7 @@ def test_memref_transforms_with_transpose(self):
       ty_in = ir.MemRefType.get(
           (64, 128),
           ir.BF16Type.get(),
-          memory_space=ir.Attribute.parse("#gpu.address_space<workgroup>"),
+          memory_space=mgpu_utils.smem(),
       )
       ref = memref.alloc(ty_in, [], [])
 

tests/mosaic/gpu_test.py

@@ -39,6 +39,7 @@
 from jax.experimental.mosaic.gpu import dialect as mgpu_dialect  # pylint: disable=g-importing-member
 from jax.experimental.mosaic.gpu import fragmented_array as fa
 from jax.experimental.mosaic.gpu import tcgen05
+from jax.experimental.mosaic.gpu import utils as mgpu_utils
 import jax.numpy as jnp
 import numpy as np
 try:
@@ -130,7 +131,6 @@ def copy(src: ir.Value, dst: ir.Value, swizzle: int | None = None):
     packing = 8 // bw
     if shape[-1] % packing:
       raise NotImplementedError
-    workgroup_mem = ir.Attribute.parse("#gpu.address_space<workgroup>")
     shape = (*shape[:-1], shape[-1] // packing)
     contig_strides = get_contiguous_strides(shape)
     def bitcast(ref):
@@ -145,12 +145,7 @@ def bitcast(ref):
           ir.StridedLayoutAttr.get(0, new_strides),
           ref_ty.memory_space,
       )
-      ptr_space = (
-          3
-          if ref_ty.memory_space is not None
-          and ref_ty.memory_space == workgroup_mem
-          else None
-      )
+      ptr_space = 3 if mgpu_utils.is_smem_ref(ref_ty) else None
       return ptr_as_memref(
           # NOTE: memref_ptr applies the offset in case there was any.
           memref_ptr(ref, memory_space=ptr_space),