Repository: jax-ml/jax
Pull request #26240 (2025-01-31 17:20:59)

[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout
[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout



Patch:
diff --git a/jax/_src/pallas/mosaic_gpu/primitives.py b/jax/_src/pallas/mosaic_gpu/primitives.py
index 428e2925f..7f26f5d2b 100644
--- a/jax/_src/pallas/mosaic_gpu/primitives.py
+++ b/jax/_src/pallas/mosaic_gpu/primitives.py
@@ -737,16 +737,32 @@ def _wgmma_accumulator_deref_lowering(ctx: lowering.LoweringRuleContext, acc):
 
 class Layout(enum.Enum):
   #: [m, n] matrix, where m % 64 == 0 == n % 8.
-  WGMMA = mgpu.WGMMAFragLayout
+  WGMMA = enum.auto()
   #: [m] matrix, where m % 64 == 0.
-  WGMMA_ROW = mgpu.WGMMARowFragLayout
+  WGMMA_ROW = enum.auto()
 
-  WG_SPLAT = mgpu.WGSplatFragLayout
-  WG_STRIDED = mgpu.WGStridedFragLayout
+  WG_SPLAT = enum.auto()
+  WG_STRIDED = enum.auto()
 
   def __call__(self, *args, **kwargs) -> ParameterizedLayout:
     return ParameterizedLayout(self, args, kwargs)
 
+  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
+    def check_no_args():
+      if args or kwargs:
+        raise ValueError(f"Can't instantiate {self} with arguments.")
+
+    match self:
+      case Layout.WGMMA:
+        check_no_args()
+        return mgpu.WGMMA_LAYOUT
+      case Layout.WGMMA_ROW:
+        check_no_args()
+        return mgpu.WGMMA_ROW_LAYOUT
+      case Layout.WG_SPLAT:
+        return mgpu.WGSplatFragLayout(*args, **kwargs)  # pytype: disable=missing-parameter
+      case Layout.WG_STRIDED:
+        return mgpu.WGStridedFragLayout(*args, **kwargs)
 
 @dataclasses.dataclass(frozen=True)
 class ParameterizedLayout:
@@ -754,16 +770,9 @@ class ParameterizedLayout:
   args: Sequence[Any]
   kwargs: Any
 
+  def to_mgpu(self) -> mgpu.FragmentedLayout:
+    return self.layout_cls.to_mgpu(*self.args, **self.kwargs)
 
-def _get_mgpu_layout(layout: Layout | ParameterizedLayout
-                     ) -> mgpu.FragmentedLayout:
-  if isinstance(layout, Layout):
-    return layout.value()
-  elif isinstance(layout, ParameterizedLayout):
-    return layout.layout_cls.value(*layout.args,
-                                   **layout.kwargs)
-  else:
-    raise TypeError(f"Unsupported layout: {layout}")
 
 layout_cast_p = jax_core.Primitive("layout_cast")
 
@@ -777,7 +786,7 @@ def _layout_cast_abstract_eval(x, new_layout):
 @lowering.register_lowering_rule(layout_cast_p, mgpu.ThreadSemantics.Lane)
 def _layout_cast_lowering(ctx: lowering.LoweringRuleContext, x, *, new_layout):
   del ctx  # Unused.
-  return x.to_layout(_get_mgpu_layout(new_layout))
+  return x.to_layout(new_layout.to_mgpu())
 
 
 def layout_cast(x: Any, new_layout: Layout | ParameterizedLayout):
@@ -860,7 +869,7 @@ def _broadcasted_iota_lowering(
   return mgpu.FragmentedArray.splat(
       llvm_dialect.mlir_undef(mlir_dtype),
       shape,
-      _get_mgpu_layout(layout),
+      layout.to_mgpu(),
       is_signed=is_signed,
   ).foreach(
       lambda _, idx: cast(idx[dimension]),
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index 49edda957..c485fd3a7 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -48,10 +48,8 @@ from .layout_inference import (
 from .fragmented_array import (
     FragmentedArray as FragmentedArray,
     FragmentedLayout as FragmentedLayout,
-    TILED_LAYOUT_WGMMA as TILED_LAYOUT_WGMMA,
     WGMMA_LAYOUT as WGMMA_LAYOUT,
     WGMMA_ROW_LAYOUT as WGMMA_ROW_LAYOUT,
-    WGMMAFragLayout as WGMMAFragLayout,
     WGMMARowFragLayout as WGMMARowFragLayout,
     WGSplatFragLayout as WGSplatFragLayout,
     WGStridedFragLayout as WGStridedFragLayout,
diff --git a/jax/experimental/mosaic/gpu/dialect_lowering.py b/jax/experimental/mosaic/gpu/dialect_lowering.py
index 55e8c5583..c9ddce106 100644
--- a/jax/experimental/mosaic/gpu/dialect_lowering.py
+++ b/jax/experimental/mosaic/gpu/dialect_lowering.py
@@ -259,7 +259,7 @@ def _vector_load_op_lowering_rule(
         is_signed=is_signed,
         vec_size=strided_layout.vec_size,
     )
-  elif layouts.from_layout_attr(out_layout_attr) == fa.TILED_LAYOUT_WGMMA:
+  elif layouts.from_layout_attr(out_layout_attr) == fa.WGMMA_LAYOUT:
     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)
@@ -267,7 +267,7 @@ def _vector_load_op_lowering_rule(
         transformed_ref,
         swizzle=swizzle,
         is_signed=is_signed,
-        layout=fa.TILED_LAYOUT_WGMMA,
+        layout=fa.WGMMA_LAYOUT,
     )
   else:
     raise ValueError(
@@ -664,7 +664,7 @@ def _mgpu_wgmma_op_lowering_rule(
       *inference_utils.out_layouts(wgmma_op),
   )
   is_supported_layout = (
-      lambda l: layouts.from_tiled_layout_attr(l) == fa.TILED_LAYOUT_WGMMA
+      lambda l: layouts.from_tiled_layout_attr(l) == fa.WGMMA_LAYOUT
   )
   if not all(map(is_supported_layout, fa_layouts)):
     raise ValueError("Layout mismatch")
@@ -701,7 +701,7 @@ def _mgpu_wgmma_op_lowering_rule(
 
   return [
       _fragmented_array_to_ir(
-          new_acc.value.to_layout(fa.TILED_LAYOUT_WGMMA),
+          new_acc.value.to_layout(fa.WGMMA_LAYOUT),
           wgmma_op.accumulator.type,
       )
   ]
diff --git a/jax/experimental/mosaic/gpu/fragmented_array.py b/jax/experimental/mosaic/gpu/fragmented_array.py
index 5cacab511..2cdaaad6a 100644
--- a/jax/experimental/mosaic/gpu/fragmented_array.py
+++ b/jax/experimental/mosaic/gpu/fragmented_array.py
@@ -97,7 +97,9 @@ class Tiling:
   def untile_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
     """Computes the shape of an array before tiling from its tiled shape."""
     def fail():
-      raise ValueError("Shape does not look like it's been tiled?")
+      raise ValueError(
+          f"shape {shape} is not a valid result of applying tiling {self}."
+      )
     for tile in reversed(self.tiles):
       if len(tile) > len(shape):
         fail()
@@ -372,14 +374,12 @@ def _tiled_wgmma_layout(shape: tuple[int, ...]):
 
   The tiled layout is equivalent to one described here in PTX documentation:
   https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-d
-
-  This tiled layout is equivalent to WGMMAFragLayout and will subsume it.
   """
   if len(shape) != 2:
     raise ValueError(f"Shape {shape} is not 2D")
   if shape[0] % 64 != 0 or shape[1] % 8 != 0:
     raise ValueError(f"Shape {shape} is not a multiple of 64x8")
-  return TILED_LAYOUT_WGMMA
+  return WGMMA_LAYOUT
 
 
 def _tiled_wgmma_layout_for_upcast(shape: tuple[int, ...]):
@@ -397,33 +397,6 @@ def _tiled_wgmma_layout_for_upcast(shape: tuple[int, ...]):
   )
 
 
-@dataclasses.dataclass(frozen=True)
-class WGMMAFragLayout:
-  """[m, n] matrix, where m % 64 == 0 == n % 8."""
-
-  def thread_idxs(self, shape):
-    index = ir.IndexType.get()
-    assert shape[0] % 64 == 0 and shape[1] % 8 == 0
-    tid = arith.index_cast(ir.IndexType.get(), mgpu.thread_idx())
-    tid_wg = arith.remui(tid, c(WARPGROUP_SIZE, index))
-    warp_idx = arith.divui(tid_wg, c(32, index))
-    tid_warp = arith.remui(tid_wg, c(32, index))
-    col_base = arith.muli(arith.remui(tid_warp, c(4, index)), c(2, index))
-    row_base = arith.addi(
-        arith.divui(tid_warp, c(4, index)), arith.muli(warp_idx, c(16, index))
-    )
-    for row_group in range(0, shape[0], 64):
-      for col_group in range(0, shape[1], 8):
-        for row_subgroup in range(0, 16, 8):
-          row = arith.addi(row_base, c(row_group + row_subgroup, index))
-          yield row, arith.addi(col_base, c(col_group, index))
-
-  def registers_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
-    assert len(shape) == 2
-    assert shape[0] % 64 == 0 and shape[1] % 8 == 0
-    return (shape[0] // 64, shape[1] // 8, 2, 1)
-
-
 @dataclasses.dataclass(frozen=True)
 class WGMMARowFragLayout:
   """[m] matrix, where m % 64 == 0."""
@@ -525,16 +498,14 @@ class WGStridedFragLayout:
       yield arith.addi(off, c(i * WARPGROUP_SIZE * self.vec_size, tidx.type))
 
 
-FragmentedLayout = WGSplatFragLayout | WGStridedFragLayout | WGMMAFragLayout | WGMMARowFragLayout | TiledLayout
+FragmentedLayout = WGSplatFragLayout | WGStridedFragLayout | WGMMARowFragLayout | TiledLayout
 
 
-WGMMA_LAYOUT = WGMMAFragLayout()
 WGMMA_ROW_LAYOUT = WGMMARowFragLayout()
 
 # The tiled layout is equivalent to one described here in PTX documentation:
 # https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-d
-# This tiled layout is equivalent to WGMMAFragLayout and will subsume it.
-TILED_LAYOUT_WGMMA = TiledLayout(
+WGMMA_LAYOUT = TiledLayout(
     Tiling(((64, 8), (16, 8), (8, 8), (1, 2))),
     warp_dim=-8,
     lane_dims=(-4, -3),
@@ -600,12 +571,6 @@ class FragmentedArray:
       )
 
     match self.layout:
-      # Registers are [m_tiles, n_tiles, 2 rows, 1 cols] in WGMMA layout
-      # Each element is a vector<2xdtype>
-      case WGMMAFragLayout():
-        if _registers.ndim != 4 or _registers.shape[2:] != (2, 1):
-          raise ValueError(f"Invalid register array shape: {_registers.shape}")
-
       # Registers are [m_tiles, 2 rows] in WGMMA_ROW layout
       # Each element is a dtype scalar
       case WGMMARowFragLayout():
@@ -680,16 +645,6 @@ class FragmentedArray:
               "WGMMARowFragLayout requires shape[0] to be a multiple of 64"
           )
         reg_shape = (shape[0] // 64, 2)
-      case WGMMAFragLayout():
-        if len(shape) != 2:
-          raise ValueError("WGMMAFragLayout requires a 2D shape")
-        if shape[0] % 64 or shape[1] % 8:
-          raise ValueError(
-              "WGMMAFragLayout requires shape[0] to be a multiple of 64, and"
-              " shape[1] to be a multiple of 8"
-          )
-        reg_shape = (shape[0] // 64, shape[1] // 8, 2, 1)
-        value = vector.splat(ir.VectorType.get((2,), value.type), value)
       case WGStridedFragLayout(vec_size=vec_size):
         assert shape == layout.shape
         elems = np.prod(shape)
@@ -713,9 +668,6 @@ class FragmentedArray:
   @property
   def shape(self):
     match self.layout:
-      case WGMMAFragLayout():
-        row_tiles, col_tiles = self.registers.shape[:2]
-        return (row_tiles * 64, col_tiles * 8)
       case WGMMARowFragLayout():
         row_tiles = self.registers.shape[0]
         return (row_tiles * 64,)
@@ -732,7 +684,7 @@ class FragmentedArray:
   def mlir_dtype(self):
     reg_ty = self.registers.flat[0].type
     match self.layout:
-      case WGMMAFragLayout() | WGStridedFragLayout() | TiledLayout():
+      case WGStridedFragLayout() | TiledLayout():
         return ir.VectorType(reg_ty).element_type
       case WGMMARowFragLayout() | WGSplatFragLayout():
         return reg_ty
@@ -749,7 +701,7 @@ class FragmentedArray:
       return self
     shape = self.shape
     if (
-        self.layout == TILED_LAYOUT_WGMMA
+        self.layout == WGMMA_LAYOUT
         and new_layout == WGMMA_TRANSPOSED_LAYOUT
         and utils.bitwidth(self.mlir_dtype) == 16
     ):
@@ -773,16 +725,6 @@ class FragmentedArray:
           _layout=new_layout,
           _is_signed=self.is_signed,
       )
-    if len(shape) == 2 and shape[0] % 64 == 0 and shape[1] % 8 == 0:
-      tiled_layout = _tiled_wgmma_layout(shape)
-      if (self.layout == WGMMA_LAYOUT and new_layout == tiled_layout) or (
-          self.layout == tiled_layout and new_layout == WGMMA_LAYOUT
-      ):
-        return FragmentedArray(
-            _registers=self.registers.reshape(new_layout.registers_shape(shape)),
-            _layout=new_layout,
-            _is_signed=self.is_signed,
-        )
     if not isinstance(self.layout, WGSplatFragLayout):
       raise NotImplementedError(
           f"Cannot convert from {self.layout} to {new_layout}"
@@ -1193,7 +1135,7 @@ class FragmentedArray:
     )
 
   def __getitem__(self, idx):
-    if self.layout not in (WGMMA_LAYOUT, TILED_LAYOUT_WGMMA):
+    if self.layout !=  WGMMA_LAYOUT:
       raise NotImplementedError("Only WGMMA layouts support slicing")
     base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape)
     if any(is_squeezed):
@@ -1346,7 +1288,7 @@ class FragmentedArray:
           )
         larger_ty = ir.F32Type.get()
         match self.layout:
-          case WGMMAFragLayout() | WGStridedFragLayout() | TiledLayout():
+          case WGStridedFragLayout() | TiledLayout():
             shape = ir.VectorType(self.registers.flat[0].type).shape
             upcast_ty = ir.VectorType.get(shape, larger_ty)
           case WGMMARowFragLayout() | WGSplatFragLayout():
@@ -1371,7 +1313,7 @@ class FragmentedArray:
       raise NotImplementedError(f"Unsupported conversion {cur_dtype} -> {new_dtype}")
     new_registers = np.empty_like(self.registers)
     match self.layout:
-      case WGMMAFragLayout() | WGStridedFragLayout() | TiledLayout():
+      case WGStridedFragLayout() | TiledLayout():
         shape = ir.VectorType(self.registers.flat[0].type).shape
         new_reg_ty = ir.VectorType.get(shape, new_dtype)
       case WGMMARowFragLayout() | WGSplatFragLayout():
@@ -1471,19 +1413,29 @@ class FragmentedArray:
       raise NotImplementedError
     index = ir.IndexType.get()
     i32 = ir.IntegerType.get_signless(32)
-    new_regs = np.empty(self.registers.shape[::2], dtype=object)
+    row_tile_dim = self.registers.shape[0]
+    row_subtile_dim = self.registers.shape[4]
+    new_regs = np.empty((row_tile_dim, row_subtile_dim), dtype=object)
     assert self.registers.shape[-1] == 1
     for row_tile, row_subtile in np.ndindex(new_regs.shape):
       # Reduce the registers owned by the current thread over n tiles
-      thread_result_vec = self.registers[row_tile, 0, row_subtile, 0]
+      reg_index = [0] * self.registers.ndim
+      reg_index[0] = row_tile
+      reg_index[4] = row_subtile
+      thread_result_vec = self.registers[tuple(reg_index)]
       for n_tile in range(1, self.registers.shape[1]):
+        reg_index[1] = n_tile
         thread_result_vec = op(
-            thread_result_vec, self.registers[row_tile, n_tile, row_subtile, 0]
+            thread_result_vec, self.registers[tuple(reg_index)]
         )
-      thread_result = op(
-          vector.extractelement(thread_result_vec, position=c(0, index)),
-          vector.extractelement(thread_result_vec, position=c(1, index)),
-      )
+
+      thread_result = vector.extractelement(thread_result_vec, position=c(0, index))
+      for i in range(1, self.layout.vector_length):
+        thread_result = op(
+            thread_result,
+            vector.extractelement(thread_result_vec, position=c(i, index)),
+        )
+
       # Do a shuffle to reduce in groups of 4 consecutive threads.
       result = thread_result
       for i in (1, 2):
@@ -1536,15 +1488,17 @@ class FragmentedArray:
   def broadcast_minor(self, n):
     if self.layout != WGMMA_ROW_LAYOUT:
       raise NotImplementedError
-    num_row_tiles = self.registers.shape[0]
-    num_col_tiles, rem = divmod(n, 8)
-    if rem:
+    if n % 8:
       raise ValueError("Number of columns must be divisible by 8")
-    new_regs = np.empty((num_row_tiles, num_col_tiles, 2, 1), dtype=object)
+    reg_shape = WGMMA_LAYOUT.registers_shape((self.shape[0], n))
+    new_regs = np.empty(reg_shape, dtype=object)
     dtype = self.mlir_dtype
     for (row_tile, row_subtile), reg in np.ndenumerate(self.registers):
-      new_regs[row_tile, :, row_subtile, :] = vector.splat(
-          ir.VectorType.get((2,), dtype), reg
+      tile = [slice(None)] * len(new_regs.shape)
+      tile[0] = row_tile
+      tile[4] = row_subtile
+      new_regs[tuple(tile)] = vector.splat(
+          ir.VectorType.get((WGMMA_LAYOUT.vector_length,), dtype), reg
       )
     return FragmentedArray(
         _registers=new_regs, _layout=WGMMA_LAYOUT, _is_signed=self.is_signed
@@ -1605,8 +1559,6 @@ class FragmentedArray:
         )
 
     match self.layout:
-      case WGMMAFragLayout():
-        self._store_untiled_wgmma(ref)
       case WGSplatFragLayout():
         vs_unsupported()
         self._store_untiled_splat(ref)
@@ -1653,34 +1605,6 @@ class FragmentedArray:
     for idx, reg in zip(idxs, self.registers.flat):
       vector.store(reg, ref_, idx)
 
-  def _store_untiled_wgmma(self, ref: ir.Value):
-    """Stores accumulator to a 2D memref. Not optimized at the moment."""
-    assert self.layout == WGMMA_LAYOUT
-    index = ir.IndexType.get()
-    m, n = self.shape
-    ref_ty = ir.MemRefType(ref.type)
-    if ref_ty.shape != [m, n]:
-      raise ValueError(ref.type, (m, n))
-
-    def c(x):
-      return arith.ConstantOp(index, ir.IntegerAttr.get(index, x))
-
-    tidx = arith.remui(gpu.thread_id(gpu.Dimension.x), c(WARPGROUP_SIZE))
-    lane_id = arith.remui(tidx, c(32))  # {0, 1, ..., 31}
-    warp_id = arith.divui(tidx, c(32))  # {0, 1, 2, 3}
-    row_base = arith.addi(
-        arith.divui(lane_id, c(4)), arith.muli(warp_id, c(16))
-    )
-    col_base = arith.muli(arith.remui(lane_id, c(4)), c(2))  # {0, 2, 4, 6}
-    it = np.ndenumerate(self.registers)
-    for (row_tile, col_tile, row_idx, col_zero), elem in it:
-      del col_zero
-      row = arith.addi(row_base, c(row_tile * 64 + row_idx * 8))
-      for col_idx in range(2):
-        value = vector.extractelement(elem, position=c(col_idx))
-        col = arith.addi(col_base, c(col_tile * 8 + col_idx))
-        memref.store(value, ref, [row, col])
-
   def _store_untiled_tiled(self, ref: ir.Value, *, vector_store: bool = True):
     """Stores an array with a tiled layout. Not optimized at the moment."""
     if utils.bitwidth(self.mlir_dtype) < 8:
@@ -1729,26 +1653,11 @@ class FragmentedArray:
         llvm.store(e, reg_ptr)
 
   def store_tiled(self, ref, swizzle: int | None):
-    match self.layout:
-      case WGMMAFragLayout():
-        dtype = self.mlir_dtype
-        bw = mgpu.bitwidth(dtype)
-        m, n = self.shape
-        assert m % 64 == 0  # This is implied by the layout.
-        cols_per_tile = (swizzle * 8) // bw
-        expected_shape = [m // 64, n // cols_per_tile, 64, cols_per_tile]
-        if n < cols_per_tile:  # We allow singular tiles shorter than swizzle.
-          expected_shape = [m // 64, 1, 64, cols_per_tile]
-        if ir.MemRefType(ref.type).shape != expected_shape:
-          raise ValueError(ref.type, (m, n))
-        for get, _, idxs in self.transfer_tiled(self.shape, dtype, swizzle):
-          vector.store(get(self.registers), ref, idxs)
-      case TiledLayout():
-        layout, shape = self.layout, self.shape
-        for get, _, ptr in self.transfer_tiled2(ref, swizzle, layout, shape):
-          llvm.store(get(self.registers), ptr)
-      case _:
-        raise NotImplementedError(self.layout)
+    if not isinstance(self.layout, TiledLayout):
+      raise NotImplementedError(self.layout)
+    layout, shape = self.layout, self.shape
+    for get, _, ptr in self.transfer_tiled2(ref, swizzle, layout, shape):
+      llvm.store(get(self.registers), ptr)
 
   @classmethod
   def load_tiled(
@@ -1778,20 +1687,6 @@ class FragmentedArray:
         reg_ty = ir.VectorType.get((layout.vector_length,), ref_ty.element_type)
         for _, update, ptr in cls.transfer_tiled2(ref, swizzle, layout, shape):
           update(registers, llvm.load(reg_ty, ptr))
-      case WGMMAFragLayout():
-        bw = mgpu.bitwidth(dtype)
-        m_tiles, n_tiles, m_tile_size, n_tile_size = ref_ty.shape
-        if m_tile_size != 64 or n_tile_size != ((swizzle * 8) // bw):
-          raise ValueError
-        m, n = m_tiles * m_tile_size, n_tiles * n_tile_size
-        assert m % 64 == 0  # This is implied by the layout.
-        registers = np.full(
-            (m_tiles, n // 8, 2, 1),
-            vector.splat(ir.VectorType.get((2,), dtype), c(0, dtype)),
-            dtype=object,
-        )
-        for _, update, idxs in cls.transfer_tiled((m, n), dtype, swizzle):
-          update(registers, vector.load(ir.VectorType.get((2,), dtype), ref, idxs))
       case _:
         raise NotImplementedError(layout)
     return cls(_registers=registers, _layout=layout, _is_signed=is_signed)
diff --git a/jax/experimental/mosaic/gpu/layout_inference.py b/jax/experimental/mosaic/gpu/layout_inference.py
index f0f0998b8..42128aff8 100644
--- a/jax/experimental/mosaic/gpu/layout_inference.py
+++ b/jax/experimental/mosaic/gpu/layout_inference.py
@@ -345,7 +345,7 @@ def _infer_reduction_op_layout(op: vector.ReductionOp) -> OptionalLayouts:
 
 @partial(_add_layout_inference_rule, mgpu.WGMMAOp)
 def _infer_wgmma_op_layout(wgmma_op: mgpu.WGMMAOp) -> OptionalLayouts:
-  layout = layouts_lib.to_layout_attr(fa.TILED_LAYOUT_WGMMA)
+  layout = layouts_lib.to_layout_attr(fa.WGMMA_LAYOUT)
 
   if ir.VectorType.isinstance(wgmma_op.a.type):
     return [layout, layout], [layout]
diff --git a/jax/experimental/mosaic/gpu/wgmma.py b/jax/experimental/mosaic/gpu/wgmma.py
index f3edbe639..ce0c5946a 100644
--- a/jax/experimental/mosaic/gpu/wgmma.py
+++ b/jax/experimental/mosaic/gpu/wgmma.py
@@ -16,6 +16,7 @@
 import dataclasses
 import functools
 import itertools
+import math
 
 import jax
 from jaxlib.mlir import ir
@@ -47,7 +48,7 @@ class WGMMAAccumulator:
   value: fa.FragmentedArray
 
   def __init__(self, *, _value: fa.FragmentedArray, _sync: bool = True):
-    if _value.layout not in (fa.WGMMA_LAYOUT, fa.TILED_LAYOUT_WGMMA):
+    if _value.layout != fa.WGMMA_LAYOUT:
       raise ValueError("Only WGMMA layouts supported in WGMMAAccumulator")
     self.value = _value
     if _sync:
@@ -123,7 +124,7 @@ def wgmma_m64(
     if a.mlir_dtype != ir.F16Type.get() and a.mlir_dtype != ir.BF16Type.get():
       raise ValueError(f"Unsupported A register array dtype: {a.mlir_dtype}")
     # Column count must be equal to swizzle // bytewidth.
-    if a.layout not in (fa.TILED_LAYOUT_WGMMA, fa.WGMMA_LAYOUT) or a.shape != (64, swizzle // 2):
+    if a.layout != fa.WGMMA_LAYOUT or a.shape != (64, swizzle // 2):
       raise ValueError("Unsupported A register array layout")
     if a_k_stride is not None or a_transpose is not None:
       raise ValueError("Unsupported WGMMA features with A in registers")
@@ -203,7 +204,7 @@ def wgmma_m64(
     imms += [lc(int(a_transpose)), lc(int(b_transpose))]
   elif supports_transpose:
     imms += [lc(int(b_transpose))]
-  if acc.ndim != 4 or acc.shape[0] != 1 or acc.shape[2:] != (2, 1):
+  if acc.ndim != 10 or acc.shape[0] != 1 or math.prod(acc.shape[2:]) != 2:
     raise ValueError(acc.shape)
   acc_struct_type = ir.Type.parse(
       f"!llvm.struct<({','.join(str(out_ty_field) for _ in acc_regs)})>"


Test patch:
diff --git a/tests/mosaic/gpu_layout_inference_test.py b/tests/mosaic/gpu_layout_inference_test.py
index ea83d1583..e0b8ab27c 100644
--- a/tests/mosaic/gpu_layout_inference_test.py
+++ b/tests/mosaic/gpu_layout_inference_test.py
@@ -210,7 +210,7 @@ class LayoutInferenceTest(parameterized.TestCase):
       for layout in [
           mgpu.WGSplatFragLayout(shape),
           mgpu.WGStridedFragLayout(shape, vec_size=4),
-          mgpu.TILED_LAYOUT_WGMMA,
+          mgpu.WGMMA_LAYOUT,
       ]
   )
   def test_infer_layout_from_yield_op_in_layouts_for_for_op(
@@ -278,7 +278,7 @@ class LayoutInferenceTest(parameterized.TestCase):
 
     mgpu.infer_layout(self.module)
 
-    wgmma_layout = layouts.to_layout_attr(mgpu.TILED_LAYOUT_WGMMA)
+    wgmma_layout = layouts.to_layout_attr(mgpu.WGMMA_LAYOUT)
     self.assertSequenceEqual(yield_op.attributes["in_layouts"], [wgmma_layout])
     self.assertSequenceEqual(yield_op.attributes["out_layouts"], [])
     self.assertSequenceEqual(for_op.attributes["in_layouts"], [wgmma_layout])
@@ -312,7 +312,7 @@ class LayoutInferenceTest(parameterized.TestCase):
 
   @parameterized.parameters(
       mgpu.WGStridedFragLayout((32, 4), vec_size=1),
-      mgpu.TILED_LAYOUT_WGMMA,
+      mgpu.WGMMA_LAYOUT,
   )
   def test_infer_layout_picks_non_splat_layout_over_splat_layout(
       self, layout
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index f6b94b777..56365a134 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -183,38 +183,28 @@ def copy(src: ir.Value, dst: ir.Value, swizzle: int | None = None):
   nvvm.fence_proxy(nvvm.ProxyKind.async_)
 
 
-def iota_tensor(m, n, dtype: jax.typing.DTypeLike, tiled_layout=False):
-  assert m % 64 == 0
-  assert n % 8 == 0
-  def c(i):
-    return arith.constant(index, ir.IntegerAttr.get(index, i))
+def iota_tensor(m, n, dtype):
+  """A wgmma tensor where arr[i, j] = i * N + j."""
   index = ir.IndexType.get()
-  i32 = ir.IntegerType.get_signless(32)
-  warp_id = arith.divui(gpu.thread_id(gpu.Dimension.x), c(32))
-  within_warp_id = arith.remui(gpu.thread_id(gpu.Dimension.x), c(32))
-  warp_row_start = arith.muli(warp_id, c(16))
-  within_warp_row = arith.divui(within_warp_id, c(4))
-  start_row = arith.addi(warp_row_start, within_warp_row)
-  start_col = arith.muli(arith.remui(within_warp_id, c(4)), c(2))
-  registers = np.empty((m // 64, n // 8, 2, 1), dtype=object)
-  for row_tile, col_tile, row_subtile, _ in np.ndindex(registers.shape):
-    row = arith.addi(start_row, c(row_tile * 64 + row_subtile * 8))
-    col = arith.addi(start_col, c(col_tile * 8))
-    row_value_base = arith.muli(row, c(n))
-    vec = llvm.mlir_undef(ir.VectorType.get((2,), i32))
-    for col_offset in range(2):
-      value = arith.addi(row_value_base, arith.addi(c(col_offset), col))
-      value = arith.index_cast(i32, value)
-      vec = vector.insertelement(value, vec, position=c(col_offset))
-    registers[row_tile, col_tile, row_subtile, 0] = vec
-  t = mgpu.FragmentedArray(
-      _registers=registers, _layout=mgpu.WGMMA_LAYOUT, _is_signed=True
+  mlir_dtype = utils.dtype_to_ir_type(dtype)
+  int_ty = ir.IntegerType.get_signless(bitwidth(mlir_dtype))
+  ret = mgpu.FragmentedArray.splat(
+      llvm.mlir_undef(int_ty), (m, n), is_signed=False
   )
-  if tiled_layout:
-    t = t.to_layout(mgpu.TILED_LAYOUT_WGMMA)
-  return t.astype(
-      utils.dtype_to_ir_type(dtype), is_signed=utils.is_signed(dtype)
+  ret = ret.to_layout(mgpu.WGMMA_LAYOUT)
+
+  def iota_value(_, idx):
+    assert len(idx) == 2
+    return arith.index_cast(
+        int_ty, arith.addi(idx[1], arith.muli(idx[0], c(n, index)))
+    )
+
+  ret = ret.foreach(
+      iota_value,
+      create_array=True,
+      is_signed=False,
   )
+  return ret.astype(mlir_dtype, is_signed=utils.is_signed(dtype))
 
 
 class TestCase(parameterized.TestCase):
@@ -272,13 +262,12 @@ class TestUtilTest(TestCase):
     y = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, x)(x)
     np.testing.assert_array_equal(y, x)
 
-  @parameterized.parameters(False, True)
-  def test_iota_tensor(self, tiled_layout):
+  def test_iota_tensor(self):
     m = n = 64
     def kernel(ctx, dst, _):
       f32 = ir.F32Type.get()
       index = ir.IndexType.get()
-      registers = iota_tensor(m, n, jnp.float32, tiled_layout).registers
+      registers = iota_tensor(m, n, jnp.float32).registers
       assert registers.size == 16, registers.size
       for i, vec_reg in enumerate(registers.flat):
         for j in range(2):
@@ -486,14 +475,13 @@ def get_packed_shape(strides, shape):
 class WGMMALayoutTest(TestCase):
 
   @parameterized.product(dtype=[jnp.float16, jnp.float32],
-                         tiled_layout=[False, True],
                          transposed_smem=[False, True])
-  def test_store_untiled(self, dtype, tiled_layout, transposed_smem):
+  def test_store_untiled(self, dtype, transposed_smem):
     def kernel(ctx, out, _):
       del ctx
       if transposed_smem:
         out = memref_transpose(out, (1, 0))
-      iota_tensor(64, 64, dtype, tiled_layout=tiled_layout).store_untiled(
+      iota_tensor(64, 64, dtype).store_untiled(
           out, vector_store=not transposed_smem
       )
     expected = np.arange(64 * 64, dtype=dtype).reshape(64, 64)
@@ -508,9 +496,8 @@ class WGMMALayoutTest(TestCase):
       dtype=[jnp.float32, jnp.float16, jnp.int8],
       swizzle=(32, 64, 128),
       num_col_tiles=(1, 2, 3),
-      tiled_layout=(False, True),
   )
-  def test_store_tiled(self, dtype, swizzle, num_col_tiles, tiled_layout):
+  def test_store_tiled(self, dtype, swizzle, num_col_tiles):
     mlir_dtype = utils.dtype_to_ir_type(dtype)
     if bytewidth(mlir_dtype) > 2 and swizzle == 32:
       self.skipTest("Not implemented")
@@ -520,7 +507,7 @@ class WGMMALayoutTest(TestCase):
     tiling = (64, col_tiling)
     def kernel(ctx, out, smem):
       del ctx
-      iota_tensor(m, n, dtype, tiled_layout).store_tiled(smem, swizzle=swizzle)
+      iota_tensor(m, n, dtype).store_tiled(smem, swizzle=swizzle)
       copy(smem, out, swizzle=swizzle)
     expected = (
         np.arange(m * n, dtype=dtype)
@@ -532,33 +519,6 @@ class WGMMALayoutTest(TestCase):
     )()
     np.testing.assert_array_equal(iota, expected)
 
-  @parameterized.product(
-      dtype=[jnp.float16, jnp.int8],
-      swizzle=(32, 64, 128),
-  )
-  def test_store_tiled_short_n(self, dtype, swizzle):
-    mlir_dtype = utils.dtype_to_ir_type(dtype)
-    col_tiling = swizzle // bytewidth(mlir_dtype)
-    m = 128
-    n = 16 // bytewidth(mlir_dtype)
-    tiling = (64, col_tiling)
-    def kernel(ctx, out, smem):
-      iota_tensor(m, n, dtype).store_tiled(smem, swizzle=swizzle)
-      ctx.async_copy(
-          src_ref=smem,
-          dst_ref=out,
-          swizzle=swizzle,
-          gmem_slice=(ds(0, m), ds(0, col_tiling)),
-          gmem_transform=mgpu.TileTransform(tiling),
-      )
-      ctx.await_async_copy(0)
-    smem_shape = jax.ShapeDtypeStruct((m // tiling[0], 1, *tiling), dtype)
-    expected = np.arange(m * n, dtype=dtype).reshape(m, n)
-    iota = mgpu.as_gpu_kernel(
-        kernel, (1, 1, 1), (128, 1, 1), (), expected, smem_shape
-    )()
-    np.testing.assert_array_equal(iota, expected)
-
   @parameterized.named_parameters(
       ("bf16_i8", jnp.bfloat16, jnp.int8),
       ("i8_bf16", jnp.int8, jnp.bfloat16),
@@ -792,9 +752,8 @@ class WGMMATest(TestCase):
       rhs_transpose=(False, True),
       swizzle=(32, 64, 128),
       dtype=[jnp.float16, jnp.bfloat16],
-      tiled_layout=(False, True),
   )
-  def test_wgmma_reg_lhs(self, m, n, k_steps, rhs_transpose, swizzle, dtype, tiled_layout):
+  def test_wgmma_reg_lhs(self, m, n, k_steps, rhs_transpose, swizzle, dtype):
     index = ir.IndexType.get()
 
     bytewidth = 2
@@ -817,7 +776,7 @@ class WGMMATest(TestCase):
               swizzle=swizzle,
           )
       init_acc = mgpu.WGMMAAccumulator.zero(m=m, n=n)
-      lhs_regs = iota_tensor(m, k, dtype, tiled_layout)
+      lhs_regs = iota_tensor(m, k, dtype)
       if rhs_transpose:
         rhs_smem = memref_transpose(rhs_smem, (0, 1, 3, 2))
       acc = mgpu.wgmma(init_acc, lhs_regs, rhs_smem, swizzle=swizzle)
@@ -1538,11 +1497,8 @@ class TMATest(TestCase):
     y_mut[:shape[0], :shape[1]] = 0
     np.testing.assert_array_equal(y_mut, np.zeros_like(y_mut))
 
-  @parameterized.product(
-      small_dim=(0, 1),
-      tiled_layout=(False, True),
-  )
-  def test_tma_small_tile_store(self, small_dim, tiled_layout):
+  @parameterized.product(small_dim=(0, 1),)
+  def test_tma_small_tile_store(self, small_dim):
     if small_dim == 0:
       shape = (4, 128)
     elif small_dim == 1:
@@ -1688,20 +1644,22 @@ class FragmentedArrayTest(TestCase):
   )
   def test_comparison(self, op, dtype, rhs_is_literal, m=64, n=32):
     def kernel(ctx, dst, _):
+      i8 = ir.IntegerType.get_signless(8)
       iota = iota_tensor(m, n, dtype)
       rhs = 0 if rhs_is_literal else iota + 1
-      op(iota, rhs).store_untiled(dst)
+      res = op(iota, rhs)
+      assert not res.is_signed
+      res.astype(i8, is_signed=False).store_untiled(dst)
 
-    out_shape = jax.ShapeDtypeStruct((m, n), jnp.bool)
+    out_shape = jax.ShapeDtypeStruct((m, n), jnp.int8)
     result = mgpu.as_gpu_kernel(
         kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
     )()
     iota = np.arange(m * n, dtype=dtype).reshape(m, n)
     rhs = rhs = 0 if rhs_is_literal else iota + 1
-    np.testing.assert_array_equal(result, op(iota, rhs))
+    np.testing.assert_array_equal(result, op(iota, rhs).astype(jnp.int8))
 
-  @parameterized.product(tiled_layout=(False, True))
-  def test_foreach(self, tiled_layout):
+  def test_foreach(self):
     dtype = jnp.int32
     swizzle = 128
     tile = 64, swizzle // jnp.dtype(dtype).itemsize
@@ -1716,7 +1674,7 @@ class FragmentedArrayTest(TestCase):
 
     tiling = mgpu.TileTransform(tile)
     def kernel(ctx, dst, smem):
-      x = iota_tensor(shape[0], shape[1], dtype, tiled_layout=tiled_layout)
+      x = iota_tensor(shape[0], shape[1], dtype)
       x.foreach(causal, create_array=True, is_signed=False).store_untiled(smem)
       mgpu.commit_shared()
       ctx.async_copy(src_ref=smem, dst_ref=dst)
@@ -1990,7 +1948,6 @@ class FragmentedArrayTest(TestCase):
   @parameterized.parameters(
       ([64 * 4], "WGMMA_ROW_LAYOUT"),
       ([64 * 4, 8 * 2], "WGMMA_LAYOUT"),
-      ([64 * 4, 8 * 2], "TILED_LAYOUT_WGMMA"),
   )
   def test_to_layout(self, shape, new_layout):
     def kernel(ctx, _):
@@ -2188,16 +2145,12 @@ class LayoutTest(TestCase):
     np.testing.assert_array_equal(f(), expected)
 
   @parameterized.product(
-      load_tiled=[False, True],
-      store_tiled=[False, True],
       dtype=[jnp.int8, jnp.int16, jnp.int32],
       swizzle=[16, 32, 64, 128],
       num_col_tiles=[1, 2, 3],
       row_tiling=[8, 64],
   )
-  def test_copy_tiled(self, load_tiled, store_tiled, dtype, swizzle, num_col_tiles, row_tiling):
-    if (not load_tiled or not load_tiled) and row_tiling != 64:
-      self.skipTest("Old code path does not support this")
+  def test_copy_tiled(self, dtype, swizzle, num_col_tiles, row_tiling):
     mlir_dtype = utils.dtype_to_ir_type(dtype)
     bw = bytewidth(mlir_dtype)
     col_tiling = swizzle // bw
@@ -2205,19 +2158,14 @@ class LayoutTest(TestCase):
       self.skipTest("WGMMA layout requires col_tiling % 8 == 0")
     m, n = 128, col_tiling * num_col_tiles
     tiling = (row_tiling, col_tiling)
-    tiled_layout = fa._tiled_wgmma_layout((m, n))
-    load_layout = tiled_layout if load_tiled else mgpu.TILED_LAYOUT_WGMMA
-    store_layout = tiled_layout if store_tiled else mgpu.TILED_LAYOUT_WGMMA
-    if (not load_tiled or not store_tiled) and bw == 4 and swizzle == 32:
-      self.skipTest("Old code path does not support this")
     def kernel(ctx, in_, out, smems):
       smem_in, smem_out, barrier = smems
       ctx.async_copy(src_ref=in_, dst_ref=smem_in, swizzle=swizzle, barrier=barrier)
       barrier.wait()
       t = mgpu.FragmentedArray.load_tiled(
-          smem_in, swizzle=swizzle, is_signed=True, layout=load_layout
+          smem_in, swizzle=swizzle, is_signed=True, layout=mgpu.WGMMA_LAYOUT
       )
-      t.to_layout(store_layout).store_tiled(smem_out, swizzle=swizzle)
+      t.store_tiled(smem_out, swizzle=swizzle)
       mgpu.commit_shared()
       ctx.async_copy(src_ref=smem_out, dst_ref=out, swizzle=swizzle)
       ctx.await_async_copy(0)
@@ -2248,9 +2196,8 @@ class LayoutTest(TestCase):
     expected_regs = swizzle // bytewidth(mlir_dtype) // 8
     # When the bytewidth is smaller than 2 the swizzle pattern changes every 2
     # column tiles, so we only need half the registers.
-    if load_tiled and store_tiled:  # The old code doesn't optimize properly.
-      if bytewidth(mlir_dtype) < 2:
-        expected_regs //= 2
+    if bytewidth(mlir_dtype) < 2:
+      expected_regs //= 2
     for instr in ("STS", "LDS"):
       with self.subTest(instr + " count"):
         addrs = re.findall(instr + r".* \[(.*)\]", get_sass())
@@ -2304,7 +2251,7 @@ class LayoutTest(TestCase):
       ctx.async_copy(src_ref=in_, dst_ref=smem_in, swizzle=swizzle, barrier=barrier)
       barrier.wait()
       t = mgpu.FragmentedArray.load_tiled(
-          smem_in, swizzle=swizzle, is_signed=True, layout=fa.TILED_LAYOUT_WGMMA
+          smem_in, swizzle=swizzle, is_signed=True, layout=fa.WGMMA_LAYOUT
       )
       smem_out_t = memref_transpose(smem_out, (1, 0, 3, 2))
       t.to_layout(transpose_layout).store_tiled(smem_out_t, swizzle=swizzle)