jax-ml/jax · error · NotImplementedError

f32 only supports add atomics, got {atomic}

Error message

f32 only supports add atomics, got {atomic}

What it means

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'.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3895

    elif multimem:
      assert not is_smem
      red = "multimem.red"
      scope = "sys"
      space = ".global"
      ptr_constraint = "l"
    else:
      red = "red"
      scope = "cta" if is_smem else "gpu"
      space = ".shared::cta" if is_smem else ""
      ptr_constraint = "r" if is_smem else "l"
    element_type = self.mlir_dtype
    element_bitwidth = utils.bitwidth(element_type)
    noftz = ""
    if isinstance(element_type, ir.F32Type):
      if cluster_barrier_ptr is not None:
        raise NotImplementedError("f32 not supported for async atomics")
      if atomic != "add":
        raise NotImplementedError(f"f32 only supports add atomics, got {atomic}")
      ptx_type = "f32"
    elif isinstance(element_type, ir.IntegerType) and element_bitwidth == 32:
      if atomic in ("and", "or", "xor"):
        ptx_type = "b32"
      else:
        ptx_type = "s32" if self.is_signed else "u32"
    elif isinstance(element_type, (ir.F16Type, ir.BF16Type)):
      if cluster_barrier_ptr is not None:
        raise NotImplementedError("f16/bf16 not supported for async atomics")
      if atomic not in ("add", "min", "max"):
        raise NotImplementedError(
            f"f16/bf16 only supports add, min, max atomics, got {atomic}"
        )
      if (is_smem or multimem) and atomic != "add":
        raise NotImplementedError(
            f"f16/bf16 SMEM/multimem atomics only support add, got {atomic}"
        )
      ptx_type = f"{element_type}x2"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use atomic='add' for f32
2. Change the buffer dtype to f16/bf16 (supports add/min/max) or to i32 with manual ordered-float encoding for min/max
3. Implement the min/max reduction non-atomically (e.g. via warp shuffles or a separate pass)

Example fix

// before
fa32.store_tiled_async(ref, atomic='max')
// after
# encode f32 as ordered i32 for max
fa_bits = fa32.bitcast(ir.IntegerType.get_signless(32))
fa_bits.store_tiled_async(ref_i32, atomic='max')

from jax._src.lib import ir
if isinstance(fa.mlir_dtype, ir.F32Type):
    assert atomic == 'add', 'f32 atomics only support add'

from jax._src.lib import ir

def atomic_ok_for_dtype(fa, atomic) -> bool:
    if isinstance(fa.mlir_dtype, ir.F32Type):
        return atomic == 'add'
    return True

try:
    fa.store_tiled_async(ref, atomic=atomic)
except NotImplementedError:
    if isinstance(fa.mlir_dtype, ir.F32Type):
        fa.store_tiled_async(ref, atomic='add')

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• f32 not supported for async atomics
• f16/bf16 not supported for async atomics
• f16/bf16 only supports add, min, max atomics, got {atomic}
• f16/bf16 SMEM/multimem atomics only support add, got {atomic
• Unsupported element type for atomic stores: {element_type}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8126cf6616e4937f. Report an issue: GitHub.