{"record":{"id":"8126cf6616e4937f","repo":"jax-ml/jax","slug":"f32-only-supports-add-atomics-got-atomic","errorCode":null,"errorMessage":"f32 only supports add atomics, got {atomic}","messagePattern":"f32 only supports add atomics, got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3895,"sourceCode":"    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}\"\n        )\n      ptx_type = f\"{element_type}x2\"","sourceCodeStart":3877,"sourceCodeEnd":3913,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3877-L3913","documentation":"PTX exposes red/atom for f32 only in the .add.f32 form (plus NaN-propagating variants); min/max and bitwise forms exist only for integers/half types, so store_tiled_async restricts f32 atomics to atomic='add'.","triggerScenarios":"Calling store_tiled_async with element type f32 and atomic set to anything other than 'add' — e.g. atomic='max' to track a running maximum.","commonSituations":"Reusing an integer/half max/min accumulation pattern with float32 buffers; porting kernels where atomic_max worked on f16 but the buffer was changed to f32.","solutions":["Use atomic='add' for f32","Change the buffer dtype to f16/bf16 (supports add/min/max) or to i32 with manual ordered-float encoding for min/max","Implement the min/max reduction non-atomically (e.g. via warp shuffles or a separate pass)"],"exampleFix":"// before\nfa32.store_tiled_async(ref, atomic='max')\n// after\n# encode f32 as ordered i32 for max\nfa_bits = fa32.bitcast(ir.IntegerType.get_signless(32))\nfa_bits.store_tiled_async(ref_i32, atomic='max')","handlingStrategy":"validation","validationCode":"from jax._src.lib import ir\nif isinstance(fa.mlir_dtype, ir.F32Type):\n    assert atomic == 'add', 'f32 atomics only support add'","typeGuard":"from jax._src.lib import ir\n\ndef atomic_ok_for_dtype(fa, atomic) -> bool:\n    if isinstance(fa.mlir_dtype, ir.F32Type):\n        return atomic == 'add'\n    return True","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, atomic=atomic)\nexcept NotImplementedError:\n    if isinstance(fa.mlir_dtype, ir.F32Type):\n        fa.store_tiled_async(ref, atomic='add')","preventionTips":["Encode min/max on f32 via ordered-integer bitcast instead","Keep a dtype→allowed-atomics table next to your store helpers"],"tags":["jax","mosaic","gpu","atomics","f32","not-implemented"],"backgroundTag":"unsupported-atomic-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}