{"record":{"id":"b19cb7b39933d4ad","repo":"jax-ml/jax","slug":"atomic-stores-not-supported-for-splat-layout","errorCode":null,"errorMessage":"Atomic stores not supported for splat layout","messagePattern":"Atomic stores not supported for splat layout","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3647,"sourceCode":"\n  def store_untiled(\n      self,\n      ref: ir.Value | utils.MultimemRef,\n      *,\n      swizzle: int = 16,\n      optimized: bool = True,\n      atomic: Literal[\"add\", \"max\", \"min\", \"and\", \"or\", \"xor\"] | None = None,\n  ) -> None:\n    index = ir.IndexType.get()\n    i64 = ir.IntegerType.get_signless(64)\n    if not isinstance(ref.type, ir.MemRefType):\n      raise ValueError(ref)\n    match self.layout:\n      case WGSplatFragLayout():\n        if isinstance(ref, utils.MultimemRef):\n          raise NotImplementedError(\"Splat layout does not support multimem\")\n        if atomic is not None:\n          raise NotImplementedError(\n              \"Atomic stores not supported for splat layout\"\n          )\n        # All values are the same so swizzle does not affect anything here.\n        self._store_untiled_splat(ref)\n      case WGStridedFragLayout():\n        if swizzle != 16:\n          raise ValueError(\"Only TiledLayouts support swizzling\")\n        assert isinstance(self.layout, WGStridedFragLayout)\n        vec_size = self.layout.vec_size\n        bitwidth = utils.bitwidth(self.mlir_dtype)\n        total_bits = vec_size * bitwidth\n        if total_bits % 8 != 0:\n          raise NotImplementedError(\"Vector length should be a multiple of byte size\")\n        # pyrefly: ignore[bad-argument-type]\n        for get, _update, transfer_ref, idx in self.transfer_strided(ref, vec_size):\n          if isinstance(transfer_ref, utils.MultimemRef):\n            ptr = utils.memref_ptr(utils.memref_slice(transfer_ref.ref, tuple(idx)))\n            if atomic is not None:","sourceCodeStart":3629,"sourceCodeEnd":3665,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3629-L3665","documentation":"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).","triggerScenarios":"Calling fa.store(ref, atomic='add') (or any atomic mode) on a FragmentedArray whose layout is WGSplatFragLayout.","commonSituations":"Trying to atomically accumulate a broadcast scalar coefficient into global memory; writing Hopper/Blackwell warp-specialized kernels where atomics were prototyped with splatted values.","solutions":["Convert to a layout supporting atomics first: fa.to_layout(TiledLayout(...)) or rebuild the value as a tiled/strided fragment, then store with atomic=.","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."],"exampleFix":"# before\nfa = FragmentedArray.splat(alpha)\nfa.store(ref, atomic='add')  # NotImplementedError\n\n# after\nmgpu.atomic_add(ptr, fa.registers.flat[0])  # scalar atomic\n# or: fa.to_layout(TiledLayout(...)).store(ref, atomic='add')","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu.fragmented_array import WGSplatFragLayout\nassert not (isinstance(fa.layout, WGSplatFragLayout) and atomic is not None), 'splat layout has no atomic store'","typeGuard":"def supports_atomic_store(fa):\n    return not isinstance(fa.layout, WGSplatFragLayout)","tryCatchPattern":null,"preventionTips":["Only use atomic= with TiledLayout or WGStridedFragLayout arrays.","For scalar atomic updates, emit a single atomic op on fa.registers.flat[0] instead."],"tags":["jax","mosaic-gpu","atomic","splat-layout","store"],"backgroundTag":"unsupported-layout-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}