jax-ml/jax · error · NotImplementedError

Predicate not supported for no-complete arrive

Error message

Predicate not supported for no-complete arrive

What it means

arrive with can_complete=False lowers to nvvm.mbarrier_arrive_nocomplete, which has no predicate operand in the NVVM API. Passing a predicate together with the no-complete variant raises NotImplementedError.

Source

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

        pred_ptx = "@$2"
        pred_constraint = ",b"
      count_ptx = f", {arrival_count}"
      if get_arch().major < 9:
        if arrival_count != 1:
          raise ValueError(
              "Only single-thread arrival is supported on pre-Hopper hardware"
          )
        count_ptx = ""
      llvm.inline_asm(
          ir.IntegerType.get_signless(64),
          [self.get_ptr()] + ([predicate] if predicate is not None else []),
          f"{pred_ptx} mbarrier.arrive.release.{ptx_scope}.shared::{ptx_scope}.b64 $0, [$1]{count_ptx};",
          "=l,r" + pred_constraint,
          has_side_effects=True,
      )
    else:
      if predicate is not None:
        raise NotImplementedError(
            "Predicate not supported for no-complete arrive"
        )
      count = c(arrival_count, ir.IntegerType.get_signless(32))
      nvvm.mbarrier_arrive_nocomplete(self.get_ptr(), count)

  def arrive_expect_tx(
      self, bytes: int | ir.Value, predicate: ir.Value | None = None
  ):
    if get_arch().major < 9:
      raise NotImplementedError("arrive_expect_tx is only supported on Hopper+ hardware")
    if isinstance(bytes, int):
      bytes = c(bytes, ir.IntegerType.get_signless(32))
    elif isinstance(bytes.type, ir.IndexType):
      i32 = ir.IntegerType.get_signless(32)
      bytes = arith.index_cast(i32, bytes)
    nvvm.mbarrier_arrive_expect_tx(
        self.get_ptr(), bytes, predicate=predicate, scope=self._nvvm_scope
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the predicate when can_complete=False; guard arrival upstream with when(predicate) around the call
2. Use the completing arrive (default) which supports predicates via inline asm
3. Restructure so only the designated lane calls arrive and skip predication entirely

Example fix

# before
barrier.arrive(can_complete=False, predicate=is_leader)
# after
with when(is_leader):
    barrier.arrive(can_complete=False)

assert not (predicate is not None and not can_complete), 'predicate requires completing arrive'

When it happens

Trigger: barrier.arrive(..., can_complete=False, predicate=p) — any combination of a non-None predicate and no-complete arrive.

Common situations: Adapting predicated arrival (used for single-lane arrival on Hopper) while also disabling completion counting; refactoring arrive calls so can_complete defaults changed under an existing predicate argument.

Related errors

• The cp.async implementation does not support user-defined pr
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

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