{"record":{"id":"7882ce2aa155d535","repo":"jax-ml/jax","slug":"splat-layout-does-not-support-multimem","errorCode":null,"errorMessage":"Splat layout does not support multimem","messagePattern":"Splat layout does not support multimem","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3645,"sourceCode":"      fmt_str = fmt.format(f\"[{idx_fmt}]: {{}}\")\n      utils.debug_print(fmt_str, *idx, val, uniform=False)\n\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):","sourceCodeStart":3627,"sourceCodeEnd":3663,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3627-L3663","documentation":"The WGSplatFragLayout (warpgroup splat layout, where every lane holds the same value) cannot store through MultimemRef references. Multimem loads/stores interact with the memory descriptor layout in a way splat fragments don't model, so Mosaic rejects the combination outright.","triggerScenarios":"fa.store(ref) where fa.layout is WGSplatFragLayout and ref is a utils.MultimemRef (e.g. a TMA/multimem descriptor obtained from warp-specialized memory ops).","commonSituations":"Kernels that splat a scalar (e.g. an accumulator init or alpha coefficient) and try to write it back through a TMA descriptor; adapting shared-memory store examples to multimem without changing the array's layout.","solutions":["Convert the array to a storable layout before storing: use fa.to_layout / re-fragment into a TiledLayout or WGStridedFragLayout, then store.","Store via a plain memref instead of the MultimemRef when the destination is regular global/shared memory.","If the value is truly scalar, write it with a single scalar store (e.g. via a small tiled array of one element) rather than a splat store."],"exampleFix":"# before\nfa_splat.store(multimem_ref)  # WGSplatFragLayout + MultimemRef\n\n# after\nfa_tiled = fa_splat.to_layout(TiledLayout(1, 1))\nfa_tiled.store(memref)  # plain memref, tiled layout","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu.fragmented_array import WGSplatFragLayout\nfrom jax.experimental.mosaic.gpu import utils\nif isinstance(fa.layout, WGSplatFragLayout) and isinstance(ref, utils.MultimemRef):\n    raise SystemExit('splat layout cannot store to multimem; convert layout first')","typeGuard":"def splat_multimem_conflict(fa, ref):\n    return isinstance(fa.layout, WGSplatFragLayout) and isinstance(ref, utils.MultimemRef)","tryCatchPattern":null,"preventionTips":["Reserve multimem/TMA stores for TiledLayout arrays.","Convert splat values to TiledLayout via to_layout before storing.","Keep scalar constants out of store paths; use explicit scalar writes."],"tags":["jax","mosaic-gpu","splat-layout","multimem","tma"],"backgroundTag":"unsupported-layout-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}