{"record":{"id":"215febd5cf8912e0","repo":"jax-ml/jax","slug":"f16-bf16-smem-multimem-atomics-only-support-add-g","errorCode":null,"errorMessage":"f16/bf16 SMEM/multimem atomics only support add, got {atomic}","messagePattern":"f16/bf16 SMEM/multimem atomics only support add, got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3910,"sourceCode":"      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}\"\n        )\n      ptx_type = f\"{element_type}x2\"\n      noftz = \"\" if multimem else \".noftz\"\n    else:\n      raise NotImplementedError(\n          f\"Unsupported element type for atomic stores: {element_type}\"\n      )\n    [vec_len] = vreg.type.shape\n    if element_bitwidth == 16:\n      if vec_len % 2 != 0:\n        raise NotImplementedError(\n            f\"f16/bf16 atomic stores require even vector length,\"\n            f\" got {vec_len}\"\n        )\n    i32_vec_len = vec_len * element_bitwidth // 32\n    vreg = utils.bitcast(vreg, ir.VectorType.get(\n        (i32_vec_len,), i32,","sourceCodeStart":3892,"sourceCodeEnd":3928,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3892-L3928","documentation":"For f16/bf16, only the packed .add.f16x2 red form is available for shared-memory (DSMEM) and multimem destinations; min/max exist only for global-memory single-element forms, hence the extra restriction when is_smem or multimem is set.","triggerScenarios":"Calling store_tiled_async on an f16/bf16 array with atomic='min' or 'max' while the reference targets shared memory (is_smem) or a multimem descriptor, or when the packed x2 path is required.","commonSituations":"Using min/max accumulation into a shared-memory reduction buffer at half precision inside a persistent kernel; works when the buffer is in global memory, so the smem restriction surprises developers.","solutions":["Use atomic='add' for smem/multimem f16/bf16 stores","Place the min/max buffer in global memory instead of shared memory","Accumulate min/max in f32 or i32 and narrow afterwards"],"exampleFix":"# before\nfa16.store_tiled_async(smem_ref, atomic='min')\n# after\nfa16.store_tiled_async(smem_ref, atomic='add')  # or use a gmem ref for min/max","handlingStrategy":"validation","validationCode":"from jax._src.lib import ir\nif isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)) and (is_smem or multimem):\n    assert atomic == 'add', 'smem/multimem half atomics support add only'","typeGuard":"from jax._src.lib import ir\n\ndef half_smem_atomic_ok(fa, atomic, is_smem, multimem) -> bool:\n    if not isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)):\n        return True\n    return atomic == 'add' or not (is_smem or multimem)","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, atomic=atomic)\nexcept NotImplementedError:\n    if atomic in ('min', 'max'):\n        fa.store_tiled_async(gmem_ref, atomic=atomic)  # relocate buffer","preventionTips":["Keep min/max half-precision buffers in global memory","Use add-only accumulation in shared-memory reduction buffers"],"tags":["jax","mosaic","gpu","atomics","f16","bf16","shared-memory","multimem","not-implemented"],"backgroundTag":"unsupported-atomic-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}