jax-ml/jax · error · NotImplementedError

get_global only supports semaphores, got {what}

Error message

get_global only supports semaphores, got {what}

What it means

pl.get_global (resource computation) only works for GMEM-resident semaphores: it must have memory_space GMEM and dtype pallas semaphore. Requesting anything else (e.g. an SMEM value or a regular array) raises NotImplementedError when computing kernel resources.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3716

          jaxpr,
          input_refs,
          consts,
      )

  assert len(outs) == len(jaxpr.outvars), (jaxpr, outs)
  return outs


@_register_resource_estimator(primitives.get_global_p)
def _get_global_resource_estimator(
    ctx: ResourceEstimatorContext, *, what
) -> Resources:
  if what.memory_space == gpu_core.GMEM and jnp.issubdtype(
      what.dtype, pallas_core.semaphore
  ):
    collective_axes = tuple(ctx.axis_names)
    return Resources(scoped_gmem_semaphores={collective_axes: what.size})
  raise NotImplementedError(f"get_global only supports semaphores, got {what}")


@register_lowering_rule(primitives.get_global_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(
    primitives.get_global_p, mgpu.LoweringSemantics.Warpgroup
)
def _get_global_lowering_rule(ctx: LoweringRuleContext, *, what):
  if what.memory_space == gpu_core.GMEM and jnp.issubdtype(
      what.dtype, pallas_core.semaphore
  ):
    collective_axes = tuple(ctx.module_ctx.axis_names)
    return ctx.module_ctx.reserve_semaphores(
        what.shape, collective_axes=collective_axes
    ).__enter__()
  raise NotImplementedError(f"get_global only supports semaphores, got {what}")


@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Lane)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only use pl.get_global for semaphores allocated in GMEM
2. Pass regular global arrays as normal BlockSpec'd inputs instead
3. Check value.memory_space == gpu_core.GMEM and jnp.issubdtype(dtype, pallas_core.semaphore) before calling

Example fix

// before
x = pl.get_global(ordinary_array_global)
// after
x = kernel_input_arg  # declared via in_specs
# get_global reserved for:
# sem = pl.get_global(semaphore_in_gmem)

assert what.memory_space == gpu_core.GMEM and jnp.issubdtype(what.dtype, pallas_core.semaphore)

def is_gmem_semaphore(aval) -> bool:
    return (aval.memory_space == gpu_core.GMEM
            and jnp.issubdtype(aval.dtype, pallas_core.semaphore))

When it happens

Trigger: Calling pl.get_global on a global whose dtype is not a semaphore or whose memory space is SMEM, during resource analysis of a GPU kernel.

Common situations: Trying to read arbitrary global arrays via get_global instead of block arguments; declaring a semaphore in the wrong memory space in the kernel signature.

Related errors

• Non-decrementing wait is not supported.
• Semaphore {sem_id} occurs as both fixed-id and internal.
• Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}
• DMA source semaphore must be a Ref.
• Cannot signal on a non-()-shaped semaphore: {src_sem_shape}

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