jax-ml/jax · error · NotImplementedError

Atomic stores not supported for splat layout

Error message

Atomic stores not supported for splat layout

What it means

Atomic stores are not implemented for arrays with the WGSplatFragLayout. Splat fragments hold one replicated value per warp, and there is no supported lowering that turns them into atomic memory operations (add/max/min/and/or/xor).

Source

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

  def store_untiled(
      self,
      ref: ir.Value | utils.MultimemRef,
      *,
      swizzle: int = 16,
      optimized: bool = True,
      atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None,
  ) -> None:
    index = ir.IndexType.get()
    i64 = ir.IntegerType.get_signless(64)
    if not isinstance(ref.type, ir.MemRefType):
      raise ValueError(ref)
    match self.layout:
      case WGSplatFragLayout():
        if isinstance(ref, utils.MultimemRef):
          raise NotImplementedError("Splat layout does not support multimem")
        if atomic is not None:
          raise NotImplementedError(
              "Atomic stores not supported for splat layout"
          )
        # All values are the same so swizzle does not affect anything here.
        self._store_untiled_splat(ref)
      case WGStridedFragLayout():
        if swizzle != 16:
          raise ValueError("Only TiledLayouts support swizzling")
        assert isinstance(self.layout, WGStridedFragLayout)
        vec_size = self.layout.vec_size
        bitwidth = utils.bitwidth(self.mlir_dtype)
        total_bits = vec_size * bitwidth
        if total_bits % 8 != 0:
          raise NotImplementedError("Vector length should be a multiple of byte size")
        # pyrefly: ignore[bad-argument-type]
        for get, _update, transfer_ref, idx in self.transfer_strided(ref, vec_size):
          if isinstance(transfer_ref, utils.MultimemRef):
            ptr = utils.memref_ptr(utils.memref_slice(transfer_ref.ref, tuple(idx)))
            if atomic is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a layout supporting atomics first: fa.to_layout(TiledLayout(...)) or rebuild the value as a tiled/strided fragment, then store with atomic=.
2. If the goal is a single atomic update, extract the scalar (fa.registers.flat[0]) and emit one atomic op on it directly instead of a full-array store.

Example fix

# before
fa = FragmentedArray.splat(alpha)
fa.store(ref, atomic='add')  # NotImplementedError

# after
mgpu.atomic_add(ptr, fa.registers.flat[0])  # scalar atomic
# or: fa.to_layout(TiledLayout(...)).store(ref, atomic='add')

from jax.experimental.mosaic.gpu.fragmented_array import WGSplatFragLayout
assert not (isinstance(fa.layout, WGSplatFragLayout) and atomic is not None), 'splat layout has no atomic store'

def supports_atomic_store(fa):
    return not isinstance(fa.layout, WGSplatFragLayout)

When it happens

Trigger: Calling fa.store(ref, atomic='add') (or any atomic mode) on a FragmentedArray whose layout is WGSplatFragLayout.

Common situations: Trying to atomically accumulate a broadcast scalar coefficient into global memory; writing Hopper/Blackwell warp-specialized kernels where atomics were prototyped with splatted values.

Related errors

• Arrays with the splat layout can only be stored when they ha
• Can only store scalars in warp-level lowering.
• Cannot broadcast shape {self.shape} to layout {o.layout}
• Splat layout does not support multimem
• Only TiledLayouts support swizzling

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