[Mosaic GPU] Add support for replicated warp_dim parsing and a dedicated test for parsing all canonical layouts.

PiperOrigin-RevId: 745015431

Author: dimitar-asenov

jax/experimental/mosaic/gpu/layouts.py

@@ -96,7 +96,7 @@ def is_strided_fragmented_layout(attr: ir.Attribute) -> bool:
 
 _tiled_layout_attr_pattern = re.compile(
     r"^#mosaic_gpu.TiledLayout<\[(?P<tiling>.*)\],"
-    r" warp_dim\s*=\s*(?P<warp_dim>[-\d]+),"
+    r" warp_dim\s*=\s*(?P<warp_dim>.+),"
     r" lane_dims\s*=\s*\[(?P<lane_dims>.*)\],"
     r" vector_dim\s*=\s*(?P<vector_dim>[-\d]+)>$"
 )
@@ -107,22 +107,26 @@ def to_tiled_layout_attr(
 ) -> ir.Attribute:
   """Constructs a #mosaic_gpu.TiledLayout attribute from a TiledLayout."""
 
-  def _lane_dim_str(d: int | fa.Replicated) -> str:
+  def _int_or_replicated(d: int | fa.Replicated) -> str:
     if isinstance(d, fa.Replicated):
       return f"#mosaic_gpu.Replicated<times={d.times}>"
     return str(d)
 
   tile_str = lambda tile: "[" + ", ".join(str(d) for d in tile) + "]"
   tiling = "[" + ", ".join(tile_str(tile) for tile in layout.tiling.tiles) + "]"
-  lane_dims = "[" +  ",".join(_lane_dim_str(d) for d in layout.lane_dims) + "]"
+  lane_dims = (
+      "[" + ",".join(_int_or_replicated(d) for d in layout.lane_dims) + "]"
+  )
 
   return ir.Attribute.parse(
-      f"#mosaic_gpu.TiledLayout<{tiling}, warp_dim={layout.warp_dim},"
+      f"#mosaic_gpu.TiledLayout<{tiling},"
+      f" warp_dim={_int_or_replicated(layout.warp_dim)},"
       f" lane_dims={lane_dims}, vector_dim={layout.vector_dim}>"
   )
 
 
 _list_of_lists_delimiter = re.compile(r"\]\s*,\s*\[")
+_int_pattern = re.compile(r"^(?P<num>[-\d]+)(\s*:\s*\w+)?$")
 _replicated_pattern = re.compile(
     r"^#mosaic_gpu.Replicated<\s*times\s*=\s*(?P<times>\d+)\s*>\s*$"
 )
@@ -143,11 +147,14 @@ def from_tiled_layout_attr(
         f"Expected a #mosaic_gpu.TiledLayout attribute, got {attr}"
     )
 
-  def _lane_dim(lane_dim_str: str) -> int | fa.Replicated:
-    match = _replicated_pattern.fullmatch(lane_dim_str)
+  def _int_or_replicated(replicated_dim: str) -> int | fa.Replicated:
+    match = _replicated_pattern.fullmatch(replicated_dim)
     if match:
       return fa.Replicated(int(match.group("times")))
-    return int(lane_dim_str)
+    match = _int_pattern.fullmatch(replicated_dim)
+    if match:
+      return int(match.group("num"))
+    raise ValueError(f"Unexpected format for replicated dim {replicated_dim}")
 
   tiling_str = match.group("tiling")
   tile_strings = []
@@ -156,9 +163,10 @@ def _lane_dim(lane_dim_str: str) -> int | fa.Replicated:
   tiles = tuple(tuple(map(int, ts.split(","))) for ts in tile_strings)
   return fa.TiledLayout(
       tiling=fa.Tiling(tiles),
-      warp_dim=int(match.group("warp_dim")),
+      warp_dim=_int_or_replicated(match.group("warp_dim")),
       lane_dims=tuple(
-          _lane_dim(s) for s in match.group("lane_dims").split(",")
+          _int_or_replicated(s.strip())
+          for s in match.group("lane_dims").split(",")
       ),
       vector_dim=int(match.group("vector_dim")),
   )

jaxlib/mosaic/dialect/gpu/mosaic_gpu.td

@@ -161,7 +161,7 @@ def MosaicGPU_TiledLayout : AttrDef<MosaicGPU_Dialect, "TiledLayout", []> {
 
   let parameters = (ins
     "::mlir::ArrayAttr":$tiling,
-    "int":$warp_dim,
+    "::mlir::Attribute":$warp_dim,
     "::mlir::ArrayAttr":$lane_dims,
     "int":$vector_dim
   );

tests/mosaic/gpu_dialect_test.py

@@ -593,6 +593,18 @@ def test_wgmma_b_n_dim_not_equal_to_acc_n_dim(self):
     ):
       self.module.operation.verify()
 
+  def test_tiled_layout_attr_parsing(self):
+    with ir.InsertionPoint(self.module.body):
+      for layout in (
+          mgpu.WGMMA_LAYOUT,
+          mgpu.WGMMA_ROW_LAYOUT,
+          mgpu.WGMMA_COL_LAYOUT,
+          mgpu.WGMMA_TRANSPOSED_LAYOUT,
+      ):
+        attr = layouts.to_tiled_layout_attr(layout)
+        parsed_layout = layouts.from_tiled_layout_attr(attr)
+        self.assertEqual(layout, parsed_layout)
+
 
 class DialectLoweringTest(MosaicGpuTest):