[Mosaic GPU] Support inference of additional layouts for BroadcastInDimOp.

PiperOrigin-RevId: 836237946

Author: allanrenucci

jax/experimental/mosaic/gpu/layout_inference.py

@@ -209,20 +209,23 @@ def extract_constant_from_broadcast_in_dim_expression_for_hint(
   if not isinstance(e.expression, eqns.RegisterLayout):
     return None
 
-  reduced_layout = e.expression.value
-
-  wgmma_tm, wgmma_tn = fa.WGMMA_LAYOUT.base_tile_shape
-  # TODO(bchetioui): enable generators to handle TCGEN05 layout from WGMMA_COL.
-  if reduced_layout == fa.WGMMA_COL_LAYOUT and e.axes == (1,) and e.shape[0] % wgmma_tm == 0:
-    return eqns.RegisterLayout(fa.WGMMA_LAYOUT)
-
-  if reduced_layout == fa.WGMMA_ROW_LAYOUT and e.axes == (0,) and e.shape[1] % wgmma_tn == 0:
-    return eqns.RegisterLayout(fa.WGMMA_LAYOUT)
-
-  tcgen05_tm, _ = fa.TCGEN05_LAYOUT.base_tile_shape
-  if reduced_layout == fa.TCGEN05_ROW_LAYOUT and e.axes == (0,) and e.shape[0] % tcgen05_tm == 0:
-    return eqns.RegisterLayout(fa.TCGEN05_LAYOUT)
+  candidates = [
+      fa.WGMMA_LAYOUT,
+      fa.WGMMA_TRANSPOSED_LAYOUT,
+      fa.TCGEN05_LAYOUT,
+      fa.TCGEN05_TRANSPOSED_LAYOUT,
+      tcgen05.TMEM_NATIVE_LAYOUT,
+  ]
+  if e.shape[-1] % 16 == 0:
+    candidates.append(tcgen05.fa_m64_collective_layout(e.shape[-1]))
 
+  # TODO(allanrenucci): Allow returning multiple valid candidates.
+  reduction_dims = tuple(d for d in range(len(e.shape)) if d not in e.axes)
+  for candidate in candidates:
+    if len(candidate.base_tile_shape) > len(e.shape):
+      continue
+    if candidate.reduce(reduction_dims) == e.expression.value:
+      return eqns.RegisterLayout(candidate)
   return None
 
 

tests/mosaic/gpu_layout_inference_test.py

@@ -293,38 +293,39 @@ def test_infer_layout_cast_layout(self):
     self.checkInLayouts(cast, [wgmma_layout])
     self.checkOutLayouts(cast, [wgmma_layout])
 
-  @parameterized.parameters(
-      (0, mgpu.WGMMA_ROW_LAYOUT, None),
-      (1, mgpu.WGMMA_COL_LAYOUT, None),
-      (0, None, mgpu.WGMMA_LAYOUT),
-      (1, None, mgpu.WGMMA_LAYOUT),
-      (0, mgpu.TCGEN05_ROW_LAYOUT, None),
-      (0, None, mgpu.TCGEN05_LAYOUT),
-      (1, None, mgpu.TCGEN05_LAYOUT),
+  @parameterized.product(
+      layout=(
+          mtu.RegisterLayout.WGMMA,
+          mtu.RegisterLayout.TCGEN05,
+          mtu.RegisterLayout.TCGEN05_TMEM_NATIVE,
+          mtu.RegisterLayout.TCGEN05_M64_COLLECTIVE,
+      ),
+      axis=(0, 1),
+      hint_on_input=(True, False),
   )
-  def test_infer_broadcast_in_dim_layout(self, broadcast_dim, in_cast, out_cast):
+  def test_infer_broadcast_in_dim_layout(self, layout, axis, hint_on_input):
     in_shape = (128,)
     out_shape = (128, 128)
+    dtype = ir.F32Type.get()
+    out_layout = layout.to_mgpu(out_shape, dtype)
+    in_layout = out_layout.reduce((1 - axis,))
 
     with ir.InsertionPoint(self.module.body):
-      [x] = undefs(ir.VectorType.get(in_shape, ir.F32Type.get()))
-      x = layout_cast(x, in_cast) if in_cast is not None else x
-      out_type = ir.VectorType.get(out_shape, ir.F32Type.get())
-      bcast = mgpu.dialect.BroadcastInDimOp(out_type, x, [broadcast_dim])
-      if out_cast is not None:
-        layout_cast(bcast.result, out_cast)
-
-    # The tests always expect WGMMA or TCGEN05 as the out layout.
-    if out_cast == mgpu.TCGEN05_LAYOUT or in_cast == mgpu.TCGEN05_ROW_LAYOUT:
-      out_layout = mgpu.TCGEN05_LAYOUT
-    else:
-      out_layout = mgpu.WGMMA_LAYOUT
+      [x] = undefs(ir.VectorType.get(in_shape, dtype))
+      if hint_on_input:
+        x = layout_cast(x, in_layout)
+      out_type = ir.VectorType.get(out_shape, dtype)
+      bcast = mgpu.dialect.BroadcastInDimOp(out_type, x, [axis])
+      if not hint_on_input:
+        layout_cast(bcast.result, out_layout)
 
-    in_layout = out_layout.reduce((1 - broadcast_dim,))
+    if hint_on_input and axis == 1 and layout == mtu.RegisterLayout.TCGEN05:
+      # Both TCGEN05 and WGMMA are valid layout candidates. WGMMA is tried first.
+      out_layout = fa.WGMMA_LAYOUT
 
     mgpu.infer_layout(self.module)
-    self.checkInLayouts(bcast, [layouts.to_layout_attr(in_layout)])
-    self.checkOutLayouts(bcast, [layouts.to_layout_attr(out_layout)])
+    self.checkInLayouts(bcast, [in_layout])
+    self.checkOutLayouts(bcast, [out_layout])
 
   @parameterized.parameters(
       (1, mgpu.WGMMA_LAYOUT, None, None),

tests/pallas/mosaic_gpu_test.py

@@ -2345,10 +2345,6 @@ def kernel(x_ref, y_ref, smem_ref, smem_reduced_ref, barrier_ref):
     np.testing.assert_allclose(x_result, op(x, axis=axis), atol=1e-5)
 
   def _test_broadcast_in_dim_base(self, shape, layout, *, axis, hint):
-    if not hint:
-      # When the hint is not set, inference may choose incompatible layouts.
-      # TODO(bchetioui): investigate and fix.
-      self.skip_if_wg_semantics()
     assert len(shape) == 2
 
     @functools.partial(