[MGPU] Add support for broadcast on major dim in WGStridedFragLayout.

PiperOrigin-RevId: 842680640

Author: golechwierowicz

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1860,6 +1860,14 @@ def _broadcast_in_dim_lowering_rule(
   if (isinstance(x.layout, mgpu.WGSplatFragLayout) and
       broadcast_dimensions == tuple(range(rank_diff, rank_diff + x_aval.ndim))):
     return x.broadcast(shape)
+  if (
+      isinstance(x.layout, mgpu.WGStridedFragLayout)
+      and broadcast_dimensions == tuple(range(rank_diff, y_aval.ndim))
+  ):
+    new_layout = mgpu.WGStridedFragLayout(
+        shape=y_aval.shape, vec_size=x.layout.vec_size
+    )
+    return x.broadcast_in_dim(y_aval.shape, broadcast_dimensions, new_layout)
   if not isinstance(layout := x.layout, mgpu.TiledLayout):
     raise NotImplementedError(f"Unsupported layout: {x.layout}")
   if any(d1 >= d2 for d1, d2 in zip(broadcast_dimensions[:-1], broadcast_dimensions[1:])):

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -2504,6 +2504,20 @@ def broadcast_in_dim(
       return type(self).splat(
         self.registers.item(), shape, layout, is_signed=self.is_signed
       )
+    if isinstance(self.layout, WGStridedFragLayout) and isinstance(layout, WGStridedFragLayout):
+      new_dims = set(range(len(shape))) - set(source_dimensions)
+      vec_match = self.layout.vec_size == layout.vec_size
+      broadcast_dim_match = new_dims == set(range(len(new_dims)))
+      assert layout.shape == shape, (layout.shape, shape)
+      if vec_match and broadcast_dim_match:
+        return FragmentedArray(
+            _registers=np.tile(
+                self.registers,
+                np.prod(shape[:len(new_dims)]),
+            ),
+            _layout=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

@@ -3796,16 +3796,49 @@ def kernel(ctx, *args):
     )(inp)
     np.testing.assert_array_equal(result, inp)
 
-  @parameterized.parameters((128, 128), (128, 64), (64, 128))
-  def test_broadcast_major(self, m, n):
+  @parameterized.product(
+      mns=((128, 128), (128, 64), (64, 128)),
+      layout=(mtu.RegisterLayout.WG_STRIDED, mtu.RegisterLayout.WGMMA),
+  )
+  def test_broadcast_major(self, mns, layout):
+    m, n = mns
+
+    if n < 128 and layout == mtu.RegisterLayout.WG_STRIDED:
+      self.skipTest(f"{n=} < 128 not supported for {layout=}")
+
+    dtype = jnp.float16
+    load_layout = (
+        layout.to_mgpu((n,), dtype)
+        if layout == mtu.RegisterLayout.WG_STRIDED
+        else mgpu.WGMMA_COL_LAYOUT
+    )
+    broadcast_layout = (
+        mgpu.WGStridedFragLayout((m, n), load_layout.vec_size)
+        if layout == mtu.RegisterLayout.WG_STRIDED
+        else layout.to_mgpu((m, n), dtype)
+    )
+
+    def load(gmem_input):
+      match layout:
+        case mtu.RegisterLayout.WG_STRIDED:
+          return mgpu.FragmentedArray.load_strided(
+              gmem_input, vec_size=load_layout.vec_size
+          )
+        case mtu.RegisterLayout.WGMMA:
+          return mgpu.FragmentedArray.load_untiled(
+              gmem_input, layout=mgpu.WGMMA_COL_LAYOUT, optimized=False
+          )
+        case _:
+          raise NotImplementedError(f"Unsupported layout: {layout}")
+
     def kernel(ctx, gmem_input, gmem_output, _):
-      t = mgpu.FragmentedArray.load_untiled(
-          gmem_input, layout=mgpu.WGMMA_COL_LAYOUT, optimized=False
+      t = load(gmem_input)
+      t.broadcast_in_dim((m, n), (1,), broadcast_layout).store_untiled(
+          gmem_output, optimized=False
       )
-      t.broadcast_in_dim((m, n), (1,), mgpu.WGMMA_LAYOUT).store_untiled(gmem_output, optimized=False)
 
-    inp = self.prng.uniform(-1, 1, (n,)).astype(jnp.float16)
-    out_shape = jax.ShapeDtypeStruct((m, n), jnp.float16)
+    inp = self.prng.uniform(-1, 1, (n,)).astype(dtype)
+    out_shape = jax.ShapeDtypeStruct((m, n), dtype)
     result = mgpu.as_gpu_kernel(
         kernel, (1, 1, 1), (128, 1, 1), (inp,), out_shape, inp
     )(inp)

tests/pallas/mosaic_gpu_test.py

@@ -2407,6 +2407,22 @@ def test_broadcast_in_dim_does_not_crash_on_small_shape(self):
         shape, plgpu.Layout.TCGEN05_TMEM_NATIVE, axis=1, hint=False
     )
 
+  def test_broadcast_in_dim_wg_strided_majormost_dim(self):
+    self.skip_if_wg_semantics()
+    @functools.partial(
+        self.pallas_call,
+        out_shape=jax.ShapeDtypeStruct((256, 128), jnp.float32),
+    )
+    def kernel(x_ref, y_ref):
+      to_be_broadcasted = plgpu.load(
+          x_ref, (), layout=plgpu.Layout.WG_STRIDED((128,), 1)
+      )
+      broadcasted = lax.broadcast_in_dim(to_be_broadcasted, (256, 128), (1,))
+      y_ref[...] = broadcasted
+
+    result = jax.random.uniform(jax.random.key(0), shape=(128,), dtype=jnp.float32)
+    np.testing.assert_array_equal(kernel(result), jnp.broadcast_to(result[None,:], (256, 128)))
+
   def test_broadcast_in_dim_tcgen05_native_layout(self):
     @functools.partial(
         self.kernel,