Repository: jax-ml/jax
Pull request #29242 (2025-06-04 16:00:28)

[Pallas/Mosaic GPU] Expose the new `TCGEN05_COL` layout.
[Pallas/Mosaic GPU] Expose the new `TCGEN05_COL` layout.

`TCGEN05_COL` is to `TCGEN05` what `WGMMA_ROW` is to `WGMMA`.



Patch:
diff --git a/jax/_src/pallas/mosaic_gpu/primitives.py b/jax/_src/pallas/mosaic_gpu/primitives.py
index f37a003f4..41c8a673b 100644
--- a/jax/_src/pallas/mosaic_gpu/primitives.py
+++ b/jax/_src/pallas/mosaic_gpu/primitives.py
@@ -1452,6 +1452,7 @@ class Layout(enum.Enum):
 
   TCGEN05 = enum.auto()
   TCGEN05_ROW = enum.auto()
+  TCGEN05_COL = enum.auto()
 
   def __call__(self, *args, **kwargs) -> ParameterizedLayout:
     return ParameterizedLayout(self, args, kwargs)
@@ -1484,6 +1485,9 @@ class Layout(enum.Enum):
       case Layout.TCGEN05_ROW:
         check_no_args()
         return mgpu.TCGEN05_ROW_LAYOUT
+      case Layout.TCGEN05_COL:
+        check_no_args()
+        return mgpu.TCGEN05_COL_LAYOUT
 
 @dataclasses.dataclass(frozen=True)
 class ParameterizedLayout:
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index 37bfa227f..f4b65c976 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -106,4 +106,5 @@ from .wgmma import (
 from .tcgen05 import (
   LAYOUT as TCGEN05_LAYOUT,  # noqa: F401
   ROW_LAYOUT as TCGEN05_ROW_LAYOUT,  # noqa: F401
+  COL_LAYOUT as TCGEN05_COL_LAYOUT,  # noqa: F401
 )
diff --git a/jax/experimental/mosaic/gpu/fragmented_array.py b/jax/experimental/mosaic/gpu/fragmented_array.py
index 3554ed958..40d95f65f 100644
--- a/jax/experimental/mosaic/gpu/fragmented_array.py
+++ b/jax/experimental/mosaic/gpu/fragmented_array.py
@@ -461,6 +461,12 @@ class TiledLayout:
         new_vector_dim,
     )
 
+  def reduce(self, axes: Sequence[int]) -> TiledLayout:
+    reduced_layout = self
+    for a in sorted(axes, reverse=True):
+      reduced_layout = reduced_layout.remove_dimension(a)
+    return reduced_layout
+
 
 def _tiled_wgmma_layout(shape: tuple[int, ...]):
   """Returns the tiled layout relevant for WGMMA operations.
@@ -1825,9 +1831,7 @@ class FragmentedArray:
       out_reg = vector.extractelement(out_reg, position=c(0, index))
       out_regs = np.asarray(out_reg, dtype=object)
     else:
-      reduced_layout = layout
-      for a in sorted(axis, reverse=True):
-        reduced_layout = reduced_layout.remove_dimension(a)
+      reduced_layout = layout.reduce(axis)
       out_regs = out_regs.reshape(reduced_layout.registers_shape(reduced_logical_shape))
     return FragmentedArray(
         _registers=out_regs, _layout=reduced_layout, _is_signed=self.is_signed
diff --git a/jax/experimental/mosaic/gpu/tcgen05.py b/jax/experimental/mosaic/gpu/tcgen05.py
index 91797fe65..32f045314 100644
--- a/jax/experimental/mosaic/gpu/tcgen05.py
+++ b/jax/experimental/mosaic/gpu/tcgen05.py
@@ -50,6 +50,13 @@ ROW_LAYOUT = fa.TiledLayout(
     lane_dims=(-3, fa.Replicated(times=4)),
     vector_dim=-1
 )
+# COL_LAYOUT is to LAYOUT as WGMMA_COL_LAYOUT is to WGMMA_LAYOUT.
+COL_LAYOUT = fa.TiledLayout(
+    fa.Tiling(tiles=((8,), (8,), (8,), (2,))),
+    warp_dim=fa.Replicated(times=4),
+    lane_dims=(fa.Replicated(times=8), -2),
+    vector_dim=-1
+)
 # A layout resembling the logical organization of TMEM. The 128 rows in a tile
 # are assigned to 128 lanes in the warpgroup. Useful when the result needs to be
 # processed in registers and then stored back into TMEM. Should not be used if


Test patch:
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index 31262dcae..a6828a31d 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -2289,6 +2289,20 @@ class FragmentedArrayTest(TestCase):
     result = m * n if reduce_both else n
     np.testing.assert_array_equal(kernel_fn(), jnp.full((m,), result, dtype))
 
+  @parameterized.named_parameters(
+      ("wgmma_row", fa.WGMMA_LAYOUT, fa.WGMMA_ROW_LAYOUT, 1),
+      ("wgmma_col", fa.WGMMA_LAYOUT, fa.WGMMA_COL_LAYOUT, 0),
+      ("tcgen05_row", tcgen05.LAYOUT, tcgen05.ROW_LAYOUT, 1),
+      ("tcgen05_col", tcgen05.LAYOUT, tcgen05.COL_LAYOUT, 0),
+  )
+  def test_layout_reduction_definition(self, layout, expected_reduced_layout, axis):
+    def squeeze_shape(shape):
+      return tuple(s for s in shape if s != 1)
+    reduced_layout = layout.reduce((axis,))
+    tiled_shape = squeeze_shape(reduced_layout.tiled_tiling_shape)
+    expected_tiled_shape = squeeze_shape(expected_reduced_layout.tiled_tiling_shape)
+    self.assertEqual(tiled_shape, expected_tiled_shape)
+
   @parameterized.product(
       op=(arith.addf, arith.maximumf),
       m=(64, 128),