{"record":{"id":"7e76561d98284149","repo":"jax-ml/jax","slug":"barriers-are-required-for-tma-gmem-smem-copies","errorCode":null,"errorMessage":"Barriers are required for TMA GMEM -> SMEM copies","messagePattern":"Barriers are required for TMA GMEM -> SMEM copies","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1289,"sourceCode":"      collective = ()\n    if not isinstance(gmem_transform, tuple):\n      gmem_transform = (gmem_transform,)\n    if not isinstance(gmem_slice, tuple):\n      gmem_slice = (gmem_slice,)\n\n    if reduction_op is not None:\n      if implementation != AsyncCopyImplementation.TMA:\n        raise ValueError(\"Only the TMA implementation supports reductions\")\n      if not _is_tma_reduction_op_supported(reduction_op, element_type):\n        raise ValueError(\n            f\"Reduction op {reduction_op} not supported by the TMA\"\n            f\" implementation for element type {element_type}\"\n        )\n\n    if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty):\n      gmem_ref, smem_ref = src_ref, dst_ref\n      if implementation == AsyncCopyImplementation.TMA and barrier is None:\n        raise ValueError(\"Barriers are required for TMA GMEM -> SMEM copies\")\n      if arrive is None:\n        arrive = True  # Arrive by default\n    elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None:\n      gmem_ref, smem_ref = dst_ref, src_ref\n      if barrier is not None:\n        raise ValueError(\"Barriers are unsupported for SMEM -> GMEM copies\")\n      if arrive is None:\n        arrive = True  # Commit this copy to the async group by default\n    else:\n      raise ValueError(\"Only SMEM <-> GMEM copies supported\")\n\n    if collective and gmem_ref is dst_ref:\n      raise ValueError(\"Only GMEM -> SMEM copies can be collective\")\n\n    (\n        slice_shape,\n        untransformed_slice_shape,\n        dyn_base_indices,","sourceCodeStart":1271,"sourceCodeEnd":1307,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1271-L1307","documentation":"TMA loads from GMEM into SMEM are asynchronous hardware operations whose completion must be tracked by an mbarrier. Mosaic therefore requires a barrier for GMEM -> SMEM TMA copies; omitting the barrier argument raises this ValueError.","triggerScenarios":"Calling async_copy(src_gmem, dst_smem, implementation=AsyncCopyImplementation.TMA) without passing barrier=..., so the load completion cannot be synchronized.","commonSituations":"Migrating a working LDGSTS-based copy to TMA and forgetting the mbarrier; new Mosaic kernels copied from examples that omit barrier setup; consuming SMEM data immediately after the copy and hitting races instead.","solutions":["Pass a barrier (e.g. mgpu.Barrier/await barrier object) to async_copy for TMA GMEM -> SMEM loads, and wait on it before reading SMEM.","Use arrive=True (the default) so the TMA transaction signals the barrier on completion.","If you cannot provide a barrier, use a non-TMA implementation for this copy."],"exampleFix":"// before\nctx.async_copy(gmem_ref, smem_ref, ..., implementation=mgpu.AsyncCopyImplementation.TMA)\n// after\nbar = mgpu.Barrier(mgpu.MemRef(() , jnp.uint32), 1)\nctx.async_copy(gmem_ref, smem_ref, ..., barrier=bar, implementation=mgpu.AsyncCopyImplementation.TMA)\nbar.await_value(1)  # before consuming smem_ref","handlingStrategy":"validation","validationCode":"is_load = src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty)\nif is_load and implementation == mgpu.AsyncCopyImplementation.TMA:\n    assert barrier is not None, 'TMA GMEM->SMEM copies require barrier='","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(gmem_ref, smem_ref, implementation=mgpu.AsyncCopyImplementation.TMA)\nexcept ValueError as e:\n    if 'Barriers are required' in str(e):\n        bar = mgpu.Barrier(...)\n        ctx.async_copy(gmem_ref, smem_ref, barrier=bar,\n                       implementation=mgpu.AsyncCopyImplementation.TMA)\n    else:\n        raise","preventionTips":["Always create an mbarrier alongside SMEM buffers used for TMA loads.","Wait on the barrier before reading SMEM to avoid data races.","Use code-generation helpers that always thread barriers through async_copy calls."],"tags":["jax","mosaic-gpu","tma","barrier","synchronization"],"backgroundTag":"missing-synchronization-primitive","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}