jax-ml/jax · error · IndexError

Barrier offset {offset} is out of bounds

Error message

Barrier offset {offset} is out of bounds

What it means

BarrierGroup.__getitem__ validates static integer offsets against the group's barrier count; offset >= num_barriers raises IndexError. It builds a sub-BarrierRef at a byte offset into the mbarrier array.

Source

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

      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:
      for offset in range(self.num_barriers):
        yield self[offset]

  def __getitem__(self, offset: ir.Value | int) -> "BarrierRef":
    i32 = ir.IntegerType.get_signless(32)
    if isinstance(offset, int):
      if offset >= self.num_barriers:
        raise IndexError(f"Barrier offset {offset} is out of bounds")
      offset = c(offset, i32)
    elif isinstance(offset.type, ir.IndexType):
      offset = arith.index_castui(i32, offset)
    elif offset.type != i32:
      raise ValueError(f"Expected a dynamic index or an integer, got {offset}")
    return BarrierRef(
        self.base_address,
        arith.addi(self.offset, offset),
        self.phases,
        1,
    )

  @property
  def _ptx_scope(self) -> str:
    if self.base_address.type == ir.Type.parse("!llvm.ptr<7>"):
      return "cluster"
    return "cta"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use range(num_barriers) exactly; assert offset < group.num_barriers in debug
2. Increase the barrier allocation if the index is legitimately needed (subject to the 32-barrier limit)
3. Derive indices from the same constant used for allocation

Example fix

# before
for i in range(num_stages + 1):
    group[i].wait_parity(0)
# after
for i in range(num_stages):
    group[i].wait_parity(0)

assert 0 <= offset < group.num_barriers, (offset, group.num_barriers)

When it happens

Trigger: group[32] on a group initialized with 32 barriers; indexing with a Python int computed from loop induction that can reach num_barriers.

Common situations: Iterating over barriers for multi-stage pipelines with an off-by-one range (range(n+1)); reinitializing with fewer barriers but keeping old hardcoded offsets; confusing barrier index with arrival count.

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/8cf01f6348075ccf. Report an issue: GitHub.