{"record":{"id":"28269c1a8ecb755e","repo":"jax-ml/jax","slug":"f32-not-supported-for-async-atomics","errorCode":null,"errorMessage":"f32 not supported for async atomics","messagePattern":"f32 not supported for async atomics","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3893,"sourceCode":"      space = \"\"\n      ptr_constraint = \"l\"\n    elif multimem:\n      assert not is_smem\n      red = \"multimem.red\"\n      scope = \"sys\"\n      space = \".global\"\n      ptr_constraint = \"l\"\n    else:\n      red = \"red\"\n      scope = \"cta\" if is_smem else \"gpu\"\n      space = \".shared::cta\" if is_smem else \"\"\n      ptr_constraint = \"r\" if is_smem else \"l\"\n    element_type = self.mlir_dtype\n    element_bitwidth = utils.bitwidth(element_type)\n    noftz = \"\"\n    if isinstance(element_type, ir.F32Type):\n      if cluster_barrier_ptr is not None:\n        raise NotImplementedError(\"f32 not supported for async atomics\")\n      if atomic != \"add\":\n        raise NotImplementedError(f\"f32 only supports add atomics, got {atomic}\")\n      ptx_type = \"f32\"\n    elif isinstance(element_type, ir.IntegerType) and element_bitwidth == 32:\n      if atomic in (\"and\", \"or\", \"xor\"):\n        ptx_type = \"b32\"\n      else:\n        ptx_type = \"s32\" if self.is_signed else \"u32\"\n    elif isinstance(element_type, (ir.F16Type, ir.BF16Type)):\n      if cluster_barrier_ptr is not None:\n        raise NotImplementedError(\"f16/bf16 not supported for async atomics\")\n      if atomic not in (\"add\", \"min\", \"max\"):\n        raise NotImplementedError(\n            f\"f16/bf16 only supports add, min, max atomics, got {atomic}\"\n        )\n      if (is_smem or multimem) and atomic != \"add\":\n        raise NotImplementedError(\n            f\"f16/bf16 SMEM/multimem atomics only support add, got {atomic}\"","sourceCodeStart":3875,"sourceCodeEnd":3911,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3875-L3911","documentation":"For f32 atomic stores with a cluster barrier, Mosaic would need an async red/atom PTX form with ::complete_tx semantics, and no such f32 instruction exists, so combining cluster_barrier_ptr with an f32 atomic raises NotImplementedError.","triggerScenarios":"Calling store_tiled_async with atomic='add' on an f32 FragmentedArray while also passing cluster_barrier_ptr (the TMA-style async completion barrier).","commonSituations":"Wiring an async pipeline with mbarrier complete_tx tracking and then performing f32 atomic accumulation; works for other dtypes so developers assume f32 is fine.","solutions":["Drop cluster_barrier_ptr for the f32 atomic store (make it synchronous)","Use i32/f32x2-packing tricks: store f32 bits via an integer-typed view if the backend supports it","Accumulate into f16/bf16 or integer shared-memory and convert afterwards, or perform the f32 atomic with a regular atomic op outside store_tiled_async"],"exampleFix":"// before\nfa32.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=bar)\n// after\nfa32.store_tiled_async(ref, atomic='add')  # no async barrier for f32","handlingStrategy":"validation","validationCode":"from jax._src.lib import ir\nif isinstance(fa.mlir_dtype, ir.F32Type) and cluster_barrier_ptr is not None:\n    cluster_barrier_ptr = None  # f32 atomics are sync-only\nfa.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=cluster_barrier_ptr)","typeGuard":"from jax._src.lib import ir\n\ndef f32_atomic_needs_sync(fa, barrier) -> bool:\n    return isinstance(fa.mlir_dtype, ir.F32Type) and barrier is not None","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=bar)\nexcept NotImplementedError:\n    fa.store_tiled_async(ref, atomic='add')","preventionTips":["Reserve cluster_barrier_ptr for non-atomic or supported-dtype stores","Document per-dtype async-atomic support in kernel helpers"],"tags":["jax","mosaic","gpu","atomics","f32","cluster-barrier","not-implemented"],"backgroundTag":"unsupported-atomic-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}