{"record":{"id":"b59e352c92e42b21","repo":"jax-ml/jax","slug":"arriving-on-a-collective-barrier-is-not-supported","errorCode":null,"errorMessage":"Arriving on a collective barrier is not supported in a warp context","messagePattern":"Arriving on a collective barrier is not supported in a warp context","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1575,"sourceCode":"    transforms_treedef,\n):\n  transforms = transforms_treedef.unflatten(flat_transforms)\n  barrier_aval = ctx.avals_in[0]\n  assert isinstance(barrier_aval, state_types.AbstractRef)\n  base_index = _get_barrier_base_index(barrier_aval, transforms)\n  if base_index is not None:\n    barrier = barrier[base_index]\n  sem_dtype = barrier_aval.inner_aval.dtype  # pyrefly: ignore[missing-attribute]\n  orders_tensor_core = getattr(sem_dtype, \"orders_tensor_core\", False)\n\n  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    scope = mgpu_utils.ThreadSubset.WARP\n  else:\n    scope = mgpu_utils.ThreadSubset.WARPGROUP\n\n  if isinstance(barrier, mgpu.CollectiveBarrierRef):\n    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n      raise NotImplementedError(\n          \"Arriving on a collective barrier is not supported in a warp context\"\n      )\n    barrier.arrive(orders_tensor_core)\n  elif ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n    barrier.arrive(orders_tensor_core)\n  else:\n    if scope == mgpu_utils.ThreadSubset.WARP and not orders_tensor_core:\n      arrival_count = 4\n    else:\n      arrival_count = 1\n\n    pred = ctx.module_ctx.single_lane_predicate if orders_tensor_core else None\n    barrier.arrive(\n        arrival_count=arrival_count,\n        orders_tensor_core=orders_tensor_core,\n        predicate=pred,\n        scope=scope,\n    )","sourceCodeStart":1557,"sourceCodeEnd":1593,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1557-L1593","documentation":"CollectiveBarrierRef.arrive is a warpgroup-wide collective; it cannot be issued from a warp-scoped context (PrimitiveSemantics.Warp) where only one warp participates. The lowering rejects it to avoid deadlock/incorrect mbarrier state.","triggerScenarios":"Calling barrier_arrive on a CollectiveBarrierRef inside code lowered with Warp semantics — e.g. inside warp_specialize warp bodies or when the pipeline lowered to Warp-only semantics.","commonSituations":"Warp-specialized kernels where a DMA/store warp tries to arrive on a barrier that was allocated as a collective barrier; converting a warpgroup kernel to warp specialization without changing barrier allocation.","solutions":["Issue the arrive from warpgroup-level code (outside warp-specialized bodies) so semantics are Warpgroup/WGxWARP.","Use a non-collective (per-warpgroup tracked) barrier for warp-scope arrivals.","Restructure warp specialization so the collective barrier is only touched at warpgroup scope."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"# Only arrive on collective barriers from warpgroup-scope code:\nif isinstance(barrier, mgpu.CollectiveBarrierRef):\n    assert not in_warp_specialized_body, 'use a non-collective barrier inside warps'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["In warp-specialized kernels, allocate per-warp barriers for warp-scope arrivals.","Reserve collective barriers for warpgroup-scope pipeline stages."],"tags":["mosaic-gpu","pallas","barrier","warp-specialization"],"backgroundTag":"unsupported-operation-scope","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}