{"record":{"id":"1d2997c836768dcb","repo":"jax-ml/jax","slug":"f16-bf16-only-supports-add-min-max-atomics-got","errorCode":null,"errorMessage":"f16/bf16 only supports add, min, max atomics, got {atomic}","messagePattern":"f16/bf16 only supports add, min, max atomics, got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3906,"sourceCode":"    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\"\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}\"","sourceCodeStart":3888,"sourceCodeEnd":3924,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3888-L3924","documentation":"PTX provides red/atom for f16/bf16 only for .add, .min and .max (packed f16x2 forms); there are no bitwise or other arithmetic atomic forms for half types, so store_tiled_async validates the requested atomic op against that set.","triggerScenarios":"Calling store_tiled_async on an F16Type/BF16Type array with atomic='and'/'or'/'xor' (or any op outside add/min/max).","commonSituations":"Generic bitwise-reduction code reused across dtypes; changing an integer mask-accumulation buffer to bfloat16 without adjusting the atomic op.","solutions":["Switch to atomic='add'/'min'/'max' for f16/bf16 data","Use an integer buffer (i32, which supports and/or/xor as b32) for bitwise reductions","Bitcast only if semantics actually match — generally keep bitwise atomics on integer types"],"exampleFix":"# before\nfa16.store_tiled_async(ref, atomic='or')\n# after\nmask_i32 = compute_masks_as_i32()\nmask_i32.store_tiled_async(ref_i32, atomic='or')","handlingStrategy":"validation","validationCode":"from jax._src.lib import ir\nHALF_ATOMICS = ('add', 'min', 'max')\nif isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)):\n    assert atomic in HALF_ATOMICS, f'use one of {HALF_ATOMICS}'","typeGuard":"from jax._src.lib import ir\n\ndef half_atomic_supported(atomic) -> bool:\n    return atomic in ('add', 'min', 'max')","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, atomic=atomic)\nexcept NotImplementedError:\n    if atomic in ('and', 'or', 'xor'):\n        raise  # bitwise needs an int buffer; do not silently change semantics","preventionTips":["Keep bitwise atomics on integer buffers only","Validate the (dtype, atomic) pair before launching the kernel"],"tags":["jax","mosaic","gpu","atomics","f16","bf16","not-implemented"],"backgroundTag":"unsupported-atomic-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}