jax-ml/jax · error · NotImplementedError

Only up to 32 barriers per group supported

Error message

Only up to 32 barriers per group supported

What it means

BarrierGroup.initialize allocates GPU mbarrier machinery that hardware/PTX limits to at most 32 mbarriers per group (one per warp lane bit). Passing a barrier_memref with more than 32 elements raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1046

def warp_barrier():
  nvvm.bar_warp_sync(c(0xFFFFFFFF, ir.IntegerType.get_signless(32)))


@dataclasses.dataclass(frozen=True)
class BarrierRef:
  base_address: ir.Value
  offset: ir.Value
  phases: ir.Value
  num_barriers: int

  @staticmethod
  def initialize(
      barrier_memref: ir.Value, arrival_count: int = 1
  ) -> "BarrierRef":
    barrier_ty = ir.MemRefType(barrier_memref.type)
    [num_barriers] = barrier_ty.shape
    if num_barriers > 32:
      raise NotImplementedError("Only up to 32 barriers per group supported")
    i32 = ir.IntegerType.get_signless(32)
    i64 = ir.IntegerType.get_signless(64)
    address = memref_ptr(barrier_memref)
    phases = memref.alloca(ir.MemRefType.get((), i32), [], [])
    memref.store(c(0, i32), phases, [])
    predicate = single_thread_predicate(scope=ThreadSubset.BLOCK)
    for i in range(num_barriers):
      nvvm.mbarrier_init(
          getelementptr(address, [i], i64),
          c(arrival_count, i32),
          predicate=predicate,
      )
    return BarrierRef(address, c(0, i32), phases, num_barriers)

  def __iter__(self) -> Iterator["BarrierRef"]:
    if self.num_barriers == 1:
      yield self
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the barrier count to <= 32 (fewer stages, or share barriers between groups)
2. Split barriers across multiple BarrierGroups and index within each
3. Re-express synchronization so a single barrier covers multiple warps via arrival counts

Example fix

# before
bars = memref.alloca(ir.MemRefType.get((64,), i32), [], [])
group = utils.BarrierGroup.initialize(bars)
# after
bars = memref.alloca(ir.MemRefType.get((32,), i32), [], [])
group = utils.BarrierGroup.initialize(bars)

assert ir.MemRefType(bars.type).shape[0] <= 32, 'max 32 barriers per group'

When it happens

Trigger: BarrierGroup.initialize(memref with shape > 32), e.g. allocating num_stages * num_consumer_groups barriers exceeding 32, or one barrier per warp with >32 warps.

Common situations: Scaling up pipeline stages or CTA sizes in a persistent kernel; using one barrier per warpgroup with many warpgroups per CTA; porting CUTLASS-style pipelines with more barriers than a 32-bit warp mask can address.

Related errors

• Indexing barrier with {transforms} not supported in GPU inte
• Expected an `NDIndexer`, but got {transforms[0]}
• Attempting to operate on barrier without indexing, but `num_
• Expected a single barrier, got a barrier reference with shap
• copy_gmem_to_smem with a barrier is only supported Hopper an

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/bda95089301facf1. Report an issue: GitHub.