[Mosaic GPU] Add barrier transformation support to tcgen05_mma.

Also fix accumulator argument when it's dynamic.

PiperOrigin-RevId: 761713920

Author: justinjfu

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1694,6 +1694,19 @@ def convert(ty, x):
     lax.not_p: lambda ctx, x: ~x,
 })
 
+def _unary_warp_lowering_rule(impl):
+  def _lowering_rule(ctx: LoweringRuleContext, x):
+    if not all(aval_in.shape == () for aval_in in ctx.avals_in):
+      raise NotImplementedError(
+          "Non-scalar arithmetic is not supported in warp-level lowering.")
+    return impl(x)
+  return _lowering_rule
+
+mosaic_lowering_rules[gpu_core.LANExWARP_SEMANTICS].update({
+    lax.neg_p: _unary_warp_lowering_rule(lambda x: -x),
+    lax.not_p: _unary_warp_lowering_rule(lambda x: ~x)
+})
+
 mosaic_lowering_rules[gpu_core.WGxWG_SEMANTICS].update({
     lax.neg_p: _lower_fun(lambda x: jnp.subtract(0, x), multiple_results=False),
     lax.not_p: _lower_fun(
@@ -2159,6 +2172,8 @@ def _axis_index_warp_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):
 
 
 @register_lowering_rule(primitives.debug_print_p, mgpu.LoweringSemantics.Lane)
+@register_lowering_rule(primitives.debug_print_p, mgpu.LoweringSemantics.Lane,
+                        gpu_core.PrimitiveSemantics.Warp)
 def _debug_print_lowering_rule(
     ctx: LoweringRuleContext,
     *args,
@@ -2167,13 +2182,17 @@ def _debug_print_lowering_rule(
 ):
   del has_placeholders  # Unused.
   primitives.check_debug_print_format(fmt, *args)
+  scope = mgpu.ThreadSubset.WARPGROUP
+  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
+    scope = mgpu.ThreadSubset.WARP
   if not any(aval.shape for aval in ctx.avals_in):
     mgpu.debug_print(
         fmt,
         *(
             _ensure_ir_value(arg, aval.dtype)
             for arg, aval in zip(args, ctx.avals_in)
         ),
+        scope=scope
     )
   elif len(ctx.avals_in) == 1:
     [arg] = args
@@ -2451,6 +2470,8 @@ def loop(loop_index, body_args):
 
 @register_lowering_rule(lax.scan_p, mgpu.LoweringSemantics.Lane)
 @register_lowering_rule(lax.scan_p, mgpu.LoweringSemantics.Warpgroup)
+@register_lowering_rule(lax.scan_p, mgpu.LoweringSemantics.Lane,
+                        gpu_core.PrimitiveSemantics.Warp)
 def _scan_lowering_rule(
     ctx: LoweringRuleContext,
     *args,

jax/_src/pallas/mosaic_gpu/primitives.py

@@ -49,6 +49,7 @@
 import jax.numpy as jnp
 
 
+WARP_SIZE = 32
 WARPGROUP_SIZE = 128
 
 
@@ -464,7 +465,7 @@ def _copy_gmem_to_smem_lowering(
     dst_transforms_treedef,
     barrier_transforms_treedef,
     collective_axes,
-    warpgroup_sync: bool = True,
+    for_warpgroup: bool = True,
 ):
   flat_src_transforms, flat_dst_transforms, flat_barrier_transforms = (
       util.split_list(
@@ -505,15 +506,23 @@ def _copy_gmem_to_smem_lowering(
   if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
     if bytes % WARPGROUP_SIZE:
       raise NotImplementedError("Only aligned copies are supported")
-    # We arrive uniformly from each thread in the WG, so we need to divide the
-    # number of bytes by the number of threads in the WG.
-    # TODO: apaszke - Relax this. We can just select the WG leader and have it
-    # arrive with the whole transfer size, while everyone else arrives with 0.
-    # But we should continue using this scheme as it's likely to be faster.
-    bytes //= WARPGROUP_SIZE
-    if warpgroup_sync:
+    if for_warpgroup:
+      # We arrive uniformly from each thread in the WG, so we need to divide the
+      # number of bytes by the number of threads in the WG.
+      # TODO: apaszke - Relax this. We can just select the WG leader and have it
+      # arrive with the whole transfer size, while everyone else arrives with 0.
+      # But we should continue using this scheme as it's likely to be faster.
+      bytes //= WARPGROUP_SIZE
       mgpu.warpgroup_barrier()  # Make sure all reads have completed.
-    barrier.arrive_expect_tx(bytes)
+      barrier.arrive_expect_tx(bytes)
+    else:
+      # In Warp-level lowering, we arrive on each CUDA thread in a warp, but
+      # the barrier still expects a full 128 arrivals so we arrive 4 times
+      # on each CUDA thread.
+      bytes //= WARP_SIZE
+      barrier.arrive_nocomplete(3)
+      barrier.arrive_expect_tx(bytes)
+
     ctx.launch_ctx.async_copy(
         src_ref=src,
         dst_ref=dst,
@@ -549,7 +558,7 @@ def _copy_gmem_to_smem_lowering(
     copy_gmem_to_smem_p,
     mgpu.LoweringSemantics.Lane,
     primitive_semantics=gpu_core.PrimitiveSemantics.Warp,
-)(functools.partial(_copy_gmem_to_smem_lowering, warpgroup_sync=False))
+)(functools.partial(_copy_gmem_to_smem_lowering, for_warpgroup=False))
 
 
 def copy_gmem_to_smem(
@@ -713,6 +722,8 @@ def _barrier_wait_pp_eqn(
 
 
 @lowering.register_lowering_rule(barrier_wait_p, mgpu.LoweringSemantics.Lane)
+@lowering.register_lowering_rule(barrier_wait_p, mgpu.LoweringSemantics.Lane,
+                                 gpu_core.PrimitiveSemantics.Warp)
 @lowering.register_lowering_rule(barrier_wait_p, mgpu.LoweringSemantics.Warpgroup)
 def _barrier_wait_lowering(
     ctx: lowering.LoweringRuleContext,
@@ -1187,18 +1198,29 @@ def tcgen05_mma(acc: _Ref,
   else:
     b_transforms_leaves, b_transforms_tree = [], None
 
+  if isinstance(barrier, pallas_core.TransformedRef):
+    barrier_transforms_leaves, barrier_transforms_tree = jax.tree.flatten(
+        barrier.transforms)
+    barrier = barrier.ref
+  else:
+    barrier_transforms_leaves, barrier_transforms_tree = [], None
+
   tcgen05_mma_p.bind(acc, a, b, barrier, accumulate,
                       *a_transforms_leaves, *b_transforms_leaves,
+                      *barrier_transforms_leaves,
                       a_transforms_tree=a_transforms_tree,
                       b_transforms_tree=b_transforms_tree,
+                      barrier_transforms_tree=barrier_transforms_tree,
                       collective_axis=collective_axis)
 
 @tcgen05_mma_p.def_abstract_eval
 def _tcgen05_mma_abstract_eval(acc, a, b, barrier, accumulate,
                                *transforms_leaves,
                                a_transforms_tree, b_transforms_tree,
+                               barrier_transforms_tree,
                                collective_axis):
-  del (accumulate, transforms_leaves, a_transforms_tree, b_transforms_tree)
+  del (accumulate, transforms_leaves, a_transforms_tree, b_transforms_tree,
+       barrier_transforms_tree)
 
   if acc.memory_space != gpu_core.TMEM:
     raise ValueError("Accumulator must be a TMEM Ref.")
@@ -1222,6 +1244,20 @@ def _tcgen05_mma_abstract_eval(acc, a, b, barrier, accumulate,
 
   return []
 
+
+def _split_transforms(all_transforms_leaves, transforms_trees) -> list[Any]:
+  transform_leaves = []
+  for transforms_tree in transforms_trees:
+    if transforms_tree is None:
+      transform_leaves.append([])
+      continue
+    current_leaves, all_transforms_leaves = util.split_list(
+        all_transforms_leaves, [transforms_tree.num_leaves]
+    )
+    transform_leaves.append(current_leaves)
+  return transform_leaves
+
+
 @lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWG_SEMANTICS)
 @lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWARP_SEMANTICS)
 def _tcgen05_mma_lowering(
@@ -1234,16 +1270,20 @@ def _tcgen05_mma_lowering(
     *transforms_leaves,
     a_transforms_tree,
     b_transforms_tree,
+    barrier_transforms_tree,
     collective_axis,
 ):
   _, a_aval, b_aval, *_ = ctx.avals_in
   lhs_swizzle: int | None = None
+  rhs_swizzle: int | None = None
   lhs_transpose: bool = False
-  if a_transforms_tree is not None:
-    a_transforms_leaves, b_transforms_leaves = util.split_list(
-        transforms_leaves, [a_transforms_tree.num_leaves]
-    )
+  rhs_transpose: bool = False
 
+  a_transforms_leaves, b_transforms_leaves, barrier_transforms_leaves = (
+      _split_transforms(transforms_leaves,
+        [a_transforms_tree, b_transforms_tree, barrier_transforms_tree])
+  )
+  if a_transforms_tree is not None:
     a_transforms = a_transforms_tree.unflatten(a_transforms_leaves)
     a_ref, a_transforms = lowering._handle_transforms(
         ctx, a_ref, a_transforms, handle_transposes=False, handle_reshapes=True
@@ -1265,36 +1305,42 @@ def _tcgen05_mma_lowering(
     if lhs_tiling != (8, swizzle_elems):
       raise ValueError("MMA lhs tiling does not fit swizzle. "
                        f"{lhs_tiling=} expected={(8, swizzle_elems)}")
-  else:
-    b_transforms_leaves = transforms_leaves  # type: ignore
 
-  b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
-  b_ref, b_transforms = lowering._handle_transforms(
-      ctx, b_ref, b_transforms, handle_transposes=False, handle_reshapes=True
-  )
-  match b_transforms:
-    case (gpu_core.UnswizzleRef(rhs_swizzle), gpu_core.UntileRef(rhs_tiling)):
-      rhs_transpose = False
-    case (
-        gpu_core.UnswizzleRef(rhs_swizzle),
-        gpu_core.UntileRef(rhs_tiling),
-        gpu_core.TransposeRef((1, 0)),
-    ):
-      rhs_transpose = True
-    case _:
-      raise NotImplementedError(
-          f"Unsupported transforms: {b_transforms}."
-      )
+  if b_transforms_tree is not None:
+    b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
+    b_ref, b_transforms = lowering._handle_transforms(
+        ctx, b_ref, b_transforms, handle_transposes=False, handle_reshapes=True
+    )
+    match b_transforms:
+      case (gpu_core.UnswizzleRef(rhs_swizzle), gpu_core.UntileRef(rhs_tiling)):
+        rhs_transpose = False
+      case (
+          gpu_core.UnswizzleRef(rhs_swizzle),
+          gpu_core.UntileRef(rhs_tiling),
+          gpu_core.TransposeRef((1, 0)),
+      ):
+        rhs_transpose = True
+      case _:
+        raise NotImplementedError(
+            f"Unsupported transforms: {b_transforms}."
+        )
+    swizzle_elems = rhs_swizzle // b_aval.dtype.itemsize
+    if rhs_tiling != (8, swizzle_elems):
+      raise ValueError("MMA rhs tiling does not fit swizzle"
+                        f" {rhs_tiling=} expected={(8, swizzle_elems)}")
+
+  if barrier_transforms_tree is not None:
+    barrier_transforms = barrier_transforms_tree.unflatten(
+        barrier_transforms_leaves)
+    indexer = _extract_barrier_indexer(barrier_transforms)
+    if indexer is not None:
+      barrier_ref = barrier_ref.__getitem__(*map(lowering._as_index, indexer.indices))
 
-  swizzle_elems = rhs_swizzle // b_aval.dtype.itemsize
   if lhs_swizzle is None:
     lhs_swizzle = rhs_swizzle
   elif rhs_swizzle != lhs_swizzle:
     raise ValueError("MMA rhs swizzle must match lhs swizzle."
                       f" {lhs_swizzle=} {rhs_swizzle=}")
-  if rhs_tiling != (8, swizzle_elems):
-    raise ValueError("MMA rhs tiling does not fit swizzle"
-                      f" {rhs_tiling=} expected={(8, swizzle_elems)}")
   if lhs_transpose:
     if isinstance(a_ref, tcgen05.TMEMRef):
       raise ValueError("TMEM transpose not allowed.")
@@ -1303,6 +1349,9 @@ def _tcgen05_mma_lowering(
     b_ref = mgpu.memref_transpose(b_ref, (1, 0, 3, 2))
   if isinstance(accumulate, bool):
     accumulate = mgpu.c(accumulate, ir.IntegerType.get_signless(1))
+  elif isinstance(accumulate, mgpu.FragmentedArray):
+    accumulate = accumulate.registers.item()
+    assert isinstance(accumulate, ir.Value)
 
   predicate = ctx.module_ctx.single_lane_predicate
   collective = False

jax/experimental/mosaic/gpu/utils.py

@@ -144,7 +144,11 @@ def _debug_scalar_ty_format(arg):
     return "%f", arg
   raise NotImplementedError(f"Can't print the type {arg.type}")
 
-def debug_print(fmt, *args, uniform=True):
+def debug_print(fmt, *args, uniform=True, scope=None):
+  if not uniform and scope is not None:
+    raise ValueError("Cannot specify scope to a non-uniform debug_print.")
+  if scope is None:
+    scope = ThreadSubset.WARPGROUP
   type_formats = []
   new_args = []
   for arg in args:
@@ -168,7 +172,7 @@ def debug_print(fmt, *args, uniform=True):
       raise NotImplementedError(arg.type)
     type_formats.append(ty_format)
   ctx = (
-      functools.partial(single_thread, scope=ThreadSubset.WARPGROUP)
+      functools.partial(single_thread, scope=scope)
       if uniform
       else contextlib.nullcontext
   )
@@ -822,6 +826,12 @@ def arrive_expect_tx(
       bytes = arith.index_cast(i32, bytes)
     nvvm.mbarrier_arrive_expect_tx_shared(self.get_ptr(), bytes, predicate=predicate)
 
+  def arrive_nocomplete(self, count: int) -> None:
+    i64 = ir.IntegerType.get_signless(64)
+    if isinstance(count, int):
+      count = c(count, ir.IntegerType.get_signless(32))
+    nvvm.mbarrier_arrive_nocomplete_shared(i64, self.get_ptr(), count)
+
   def get_ptr(self):
     ptr = ir.Type.parse(f"!llvm.ptr<{WORKGROUP_NVPTX_ADDRESS_SPACE}>")
     i64 = ir.IntegerType.get_signless(64)