[Mosaic GPU] Add non-collective blackwell matmul example

PiperOrigin-RevId: 761718971

Author: justinjfu

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

@@ -830,6 +830,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)