[Mosaic GPU] Support FragmentedArray.broadcast_in_dim for splat to other layouts.

PiperOrigin-RevId: 842145211

Author: allanrenucci

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -1301,9 +1301,8 @@ def to_layout(self, new_layout: FragmentedLayout) -> FragmentedArray:
       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
+        self.registers.item(), self.shape, new_layout, is_signed=self.is_signed
     )
 
   def _pointwise(
@@ -2502,15 +2501,9 @@ def broadcast_in_dim(
             f" {shape[target_dim]} in shape after broadcast"
         )
     if isinstance(self.layout, WGSplatFragLayout):
-      if isinstance(layout, WGSplatFragLayout):
-        if layout.shape != shape:
-          raise ValueError(
-              f"Layout shape {layout.shape} does not match broadcast shape {shape}"
-          )
-        return FragmentedArray(
-            _registers=self.registers, _layout=layout, _is_signed=self.is_signed,
-        )
-      # TODO: Support splat to other layouts
+      return type(self).splat(
+        self.registers.item(), shape, layout, is_signed=self.is_signed
+      )
     if not isinstance(self.layout, TiledLayout) or not isinstance(layout, TiledLayout):
       raise NotImplementedError(self.layout, layout)
     if any(d1 >= d2 for d1, d2 in zip(source_dimensions, source_dimensions[1:])):

tests/mosaic/gpu_test.py

@@ -3812,6 +3812,30 @@ def kernel(ctx, gmem_input, gmem_output, _):
     out_ref = jax.lax.broadcast_in_dim(inp, (m, n), (1,))
     np.testing.assert_array_equal(result, out_ref)
 
+  @parameterized.parameters(*mtu.RegisterLayout)
+  def test_broadcast_splat(self, layout):
+    out_shape = (128, 128)
+
+    def body(ctx, out_ref, scratch):
+      del ctx, scratch
+      c42 = arith.constant(ir.IntegerType.get_signless(32), 42)
+      arr = mgpu.FragmentedArray.splat(c42, (128,), is_signed=True)
+      out_layout = layout.to_mgpu(out_shape, jnp.int32)
+      result = arr.broadcast_in_dim(out_shape, (0,), out_layout)
+      result.store_untiled(out_ref, optimized=False)
+
+    kernel = mgpu.as_gpu_kernel(
+        body,
+        grid=(1, 1, 1),
+        block=(128, 1, 1),
+        in_shape=(),
+        out_shape=jax.ShapeDtypeStruct(out_shape, jnp.int32),
+        smem_scratch_shape=[],
+    )
+    np.testing.assert_array_equal(
+        kernel(), np.full(out_shape, 42, dtype=np.int32)
+    )
+
   def test_warp_tree_reduce(self):
     def kernel(ctx, out, *_):
       del ctx