[Mosaic GPU] Test the wgmma_op lowering when a is in registers.

I had to add support for wgmma layout in vector_load. Not sure if this is useful outside the test.

PiperOrigin-RevId: 735384104

Author: dimitar-asenov

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -235,11 +235,6 @@ def _vector_load_op_lowering_rule(
       ir.ArrayAttr, vector_load_op.attributes["out_layouts"]
   )
 
-  if not layouts.is_strided_fragmented_layout(out_layout_attr):
-    raise ValueError(
-        f"{vector_load_op} has an unsupported layout: {out_layout_attr}"
-    )
-
   for i in vector_load_op.indices:
     index_defining_op = i.owner.opview
     if (
@@ -254,10 +249,29 @@ def _vector_load_op_lowering_rule(
 
   element_type = vector_load_op.result.type.element_type
   is_signed = False if ir.IntegerType.isinstance(element_type) else None
-  strided_layout = layouts.from_strided_fragmented_layout_attr(out_layout_attr)
-  fragmented_array = fa.FragmentedArray.load_strided(
-      vector_load_op.base, is_signed=is_signed, vec_size=strided_layout.vec_size
-  )
+
+  if layouts.is_strided_fragmented_layout(out_layout_attr):
+    strided_layout = layouts.from_strided_fragmented_layout_attr(
+        out_layout_attr
+    )
+    fragmented_array = fa.FragmentedArray.load_strided(
+        vector_load_op.base,
+        is_signed=is_signed,
+        vec_size=strided_layout.vec_size,
+    )
+  elif layouts.is_wgmma_fragmented_layout(out_layout_attr):
+    layout = ir.MemRefType(vector_load_op.base.type).layout
+    swizzle, transforms = memref_layout_to_swizzle_and_transforms(layout)
+    transformed_ref = transform_memref(vector_load_op.base, transforms)
+    fragmented_array = fa.FragmentedArray.load_tiled(
+        transformed_ref,
+        swizzle=swizzle,
+        is_signed=is_signed
+    )
+  else:
+    raise ValueError(
+        f"{vector_load_op} has an unsupported layout: {out_layout_attr}"
+    )
   return [_fragmented_array_to_ir(fragmented_array, vector_load_op.result.type)]
 
 

tests/mosaic/gpu_test.py

@@ -2755,6 +2755,7 @@ class TestCaseInput:
       transforms_b: tuple[Tile | Transpose | Swizzle, ...] = ()
       transpose_a: bool = False
       transpose_b: bool = False
+      load_a_in_registers: bool = False
 
     result = []
     for swizzle in [
@@ -2786,6 +2787,13 @@ class TestCaseInput:
               transforms_a=[Tile([64, k]), Swizzle(swizzle)],
               transforms_b=[Tile([k, k]), Swizzle(swizzle)],
           ),
+          TestCaseInput(
+              shape_a=[groups_m * 64, groups_k * k],
+              shape_b=[groups_k * k, groups_n * k],
+              shape_res=[groups_m * 64, groups_n * k],
+              transforms_a=[Tile([64, k]), Swizzle(swizzle)],
+              load_a_in_registers=True,
+          ),
       ])
       # The below only works for 128-byte swizzling. Regardless of transposing,
       # TMA needs the size of the last dimension to be compatible with the
@@ -2849,6 +2857,14 @@ def matmul(
       parity, _ = tma_barrier.update_parities(parities)
       mgpu_dialect.wait(dialect_barrier, parity)
 
+      # SMEM -> Registers
+      a_operand = a_smem_ref
+      zero_index = arith.constant(ir.IndexType.get(), 0)
+      if test_case.load_a_in_registers:
+        a_vector_type = ir.VectorType.get(test_case.shape_a, ab_elt_type)
+        zero_vector_indices = [zero_index] * len(test_case.shape_a)
+        a_operand = vector.load(a_vector_type, a_smem_ref, zero_vector_indices)
+
       # Computation
       shape_result = ir.MemRefType(result_gmem_ref.type).shape
       result_elt_type = ir.MemRefType(result_gmem_ref.type).element_type
@@ -2860,7 +2876,7 @@ def matmul(
       )
       result = mgpu_dialect.wgmma(
           accumulator,
-          a_smem_ref,
+          a_operand,
           b_smem_ref,
           transpose_a=test_case.transpose_a,
           transpose_b=test_case.transpose_b,
@@ -2870,8 +2886,7 @@ def matmul(
       nvvm.wgmma_wait_group_sync_aligned(0)
 
       # Registers -> SMEM
-      zero_index = arith.constant(ir.IndexType.get(), 0)
-      vector.store(result, result_smem_ref, [zero_index, zero_index])
+      vector.store(result, result_smem_ref, [zero_index] * len(shape_result))
 
       # SMEM -> GMEM
       mgpu_dialect.async_store(