{"record":{"id":"bda95089301facf1","repo":"jax-ml/jax","slug":"only-up-to-32-barriers-per-group-supported","errorCode":null,"errorMessage":"Only up to 32 barriers per group supported","messagePattern":"Only up to 32 barriers per group supported","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":1046,"sourceCode":"def warp_barrier():\n  nvvm.bar_warp_sync(c(0xFFFFFFFF, ir.IntegerType.get_signless(32)))\n\n\n@dataclasses.dataclass(frozen=True)\nclass BarrierRef:\n  base_address: ir.Value\n  offset: ir.Value\n  phases: ir.Value\n  num_barriers: int\n\n  @staticmethod\n  def initialize(\n      barrier_memref: ir.Value, arrival_count: int = 1\n  ) -> \"BarrierRef\":\n    barrier_ty = ir.MemRefType(barrier_memref.type)\n    [num_barriers] = barrier_ty.shape\n    if num_barriers > 32:\n      raise NotImplementedError(\"Only up to 32 barriers per group supported\")\n    i32 = ir.IntegerType.get_signless(32)\n    i64 = ir.IntegerType.get_signless(64)\n    address = memref_ptr(barrier_memref)\n    phases = memref.alloca(ir.MemRefType.get((), i32), [], [])\n    memref.store(c(0, i32), phases, [])\n    predicate = single_thread_predicate(scope=ThreadSubset.BLOCK)\n    for i in range(num_barriers):\n      nvvm.mbarrier_init(\n          getelementptr(address, [i], i64),\n          c(arrival_count, i32),\n          predicate=predicate,\n      )\n    return BarrierRef(address, c(0, i32), phases, num_barriers)\n\n  def __iter__(self) -> Iterator[\"BarrierRef\"]:\n    if self.num_barriers == 1:\n      yield self\n    else:","sourceCodeStart":1028,"sourceCodeEnd":1064,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L1028-L1064","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reduce the barrier count to <= 32 (fewer stages, or share barriers between groups)","Split barriers across multiple BarrierGroups and index within each","Re-express synchronization so a single barrier covers multiple warps via arrival counts"],"exampleFix":"# before\nbars = memref.alloca(ir.MemRefType.get((64,), i32), [], [])\ngroup = utils.BarrierGroup.initialize(bars)\n# after\nbars = memref.alloca(ir.MemRefType.get((32,), i32), [], [])\ngroup = utils.BarrierGroup.initialize(bars)","handlingStrategy":"validation","validationCode":"assert ir.MemRefType(bars.type).shape[0] <= 32, 'max 32 barriers per group'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate barriers as a small power-of-two count <= 32","Share barriers across warps using arrival counts instead of more barriers"],"tags":["jax","mosaic-gpu","barrier","hardware-limit","cuda"],"backgroundTag":"gpu-barrier-limit-exceeded","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}