{"record":{"id":"9de39f1ad89982dc","repo":"jax-ml/jax","slug":"copy-gmem-to-smem-without-a-barrier-is-only-suppor","errorCode":null,"errorMessage":"copy_gmem_to_smem without a barrier is only supported on pre-Hopper GPUs, which use the cp.async implementation","messagePattern":"copy_gmem_to_smem without a barrier is only supported on pre-Hopper GPUs, which use the cp\\.async implementation","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1024,"sourceCode":"          \"Only the TMA implementation supports leader_tracked copies\"\n      )\n    # cp.async does not predicate out-of-bounds accesses, so the caller has to\n    # guarantee that the copy stays in bounds.\n    if oob_mode != OOBFillMode.PROMISE_IN_BOUNDS:\n      raise ValueError(\n          \"The cp.async implementation only supports \"\n          \"oob_mode=OOBFillMode.PROMISE_IN_BOUNDS\"\n      )\n    if has_user_predicate:\n      raise NotImplementedError(\n          \"The cp.async implementation does not support user-defined predicates\"\n      )\n  else:\n    if oob_mode is None:\n      oob_mode = OOBFillMode.ZEROS\n\n    if barrier is None:\n      raise ValueError(\n          \"copy_gmem_to_smem without a barrier is only supported on pre-Hopper\"\n          \" GPUs, which use the cp.async implementation\"\n      )\n\n  i32 = ir.IntegerType.get_signless(32)\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    if (\n        ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup\n        and ctx.module_ctx.auto_barriers\n    ):\n      mgpu.warpgroup_barrier()  # Make sure all reads have completed.\n\n    if not is_cp_async:\n      assert barrier is not None\n      if bytes % WARPGROUP_SIZE:\n        raise NotImplementedError(\n            \"Only copies transferring a number of bytes divisible by the\"\n            f\" warpgroup size are supported. Got {bytes=} but warpgroup size is\"","sourceCodeStart":1006,"sourceCodeEnd":1042,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1006-L1042","documentation":"Mosaic GPU's copy_gmem_to_smem lowering requires an explicit barrier when running on Hopper+ (SM90+) GPUs, because on those architectures the copy is implemented with TMA (async_load) which needs a barrier to signal completion. Only pre-Hopper GPUs use the cp.async path where barrier=None is allowed. The check fires when barrier is None under Warpgroup lowering semantics on TMA-capable hardware.","triggerScenarios":"Calling copy_gmem_to_smem (directly or via copy_in) without passing a barrier while running on a Hopper (SM90) or newer GPU, i.e. when the lowering does not take the cp.async (is_cp_async) path.","commonSituations":"Kernels written for A100 that ran copy_gmem_to_smem(src, dst) with no barrier are run unchanged on H100; or code that conditionally omits the barrier based on an outdated device check.","solutions":["Pass a mgpu.SMEM_ALLOCATOR Barrier to copy_gmem_to_smem, e.g. allocate a barrier and pass barrier=barrier, then wait_gmem_to_smem/barrier.wait after the copy on Hopper.","If you truly want the no-barrier cp.async path, target a pre-Hopper GPU (set the Mosaic GPU arch/environment to SM80).","Upgrade kernel code to the TMA-oriented API (copy_in / device topology helpers) which manages barriers automatically."],"exampleFix":"# before\ncopy_gmem_to_smem(src_ref, dst_ref)\n# after\nbarrier = mgpu.SMEM_ALLOCATOR.get_buffer((1,), mgpu.BarrierType, barrier_init=1)\ncopy_gmem_to_smem(src_ref, dst_ref, barrier=barrier)\nbarrier.wait()","handlingStrategy":"validation","validationCode":"import jax\nfrom jax._src.pallas.mosaic_gpu import mgpu\narch = mgpu.utils.get_arch()\nneeds_barrier = arch.major >= 9\nif needs_barrier:\n    barrier = mgpu.SMEM_ALLOCATOR.get_buffer((1,), mgpu.BarrierType, barrier_init=1)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass an explicit barrier to copy_gmem_to_smem; it is valid on all architectures.","Write kernels against copy_in/higher-level helpers that pick the right mechanism per SM version."],"tags":["mosaic-gpu","pallas","gpu","barrier","tma","hopper"],"backgroundTag":"gpu-barrier-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}