[pallas:mosaic_gpu] Added FragmentedArray.to_layout

PiperOrigin-RevId: 686524192

Author: superbobry

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -232,15 +232,20 @@ def splat(cls, value, shape, layout=None, *, is_signed: bool | None = None):
     match layout:
       case WGMMARowFragLayout():
         if len(shape) != 1:
-          raise ValueError
+          raise ValueError("WGMMARowFragLayout requires a 1D shape")
         if shape[0] % 64:
-          raise ValueError
+          raise ValueError(
+              "WGMMARowFragLayout requires shape[0] to be a multiple of 64"
+          )
         reg_shape = (shape[0] // 64, 2)
       case WGMMAFragLayout():
         if len(shape) != 2:
-          raise ValueError
+          raise ValueError("WGMMAFragLayout requires a 2D shape")
         if shape[0] % 64 or shape[1] % 8:
-          raise ValueError
+          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):
@@ -283,6 +288,22 @@ def mlir_dtype(self):
       case WGMMARowFragLayout() | WGSplatFragLayout():
         return reg_ty
 
+  def to_layout(self, new_layout: FragmentedLayout):
+    """Converts the fragmented array to the given layout.
+
+    At the moment, only conversions from ``WGSplatFragLayout`` are supported.
+    """
+    if self.layout == new_layout:
+      return self
+    if not isinstance(self.layout, WGSplatFragLayout):
+      raise NotImplementedError(
+          f"Cannot convert from {self.layout} to {new_layout}"
+      )
+    [reg] = self.registers.flat
+    return type(self).splat(
+        reg, self.shape, new_layout, is_signed=self.is_signed
+    )
+
   def _pointwise(self, op, *other, output_is_signed: bool | None = None):
     is_signed = (
         output_is_signed if output_is_signed is not None else self.is_signed

tests/mosaic/gpu_test.py

@@ -1462,6 +1462,18 @@ def kernel(ctx, inp, out, smem):
     )(x)
     np.testing.assert_array_equal(result, reference)
 
+  @parameterized.parameters(
+      ([64 * 4], "WGMMA_ROW_LAYOUT"),
+      ([64 * 4, 8 * 2], "WGMMA_LAYOUT"),
+  )
+  def test_to_layout(self, shape, new_layout):
+    def kernel(ctx, _):
+      # No assertions, we are just checking there are no compile-time errors.
+      arr = mgpu.FragmentedArray.splat(c(42.0, ir.F32Type.get()), shape)
+      arr.to_layout(getattr(mgpu, new_layout))
+
+    _ = mgpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), (), (), None)()
+
 
 class ProfilerTest(TestCase):