{"record":{"id":"b410008f44db8632","repo":"jax-ml/jax","slug":"only-tiledlayouts-support-swizzling","errorCode":null,"errorMessage":"Only TiledLayouts support swizzling","messagePattern":"Only TiledLayouts support swizzling","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3654,"sourceCode":"      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:\n              self._store_register_atomic(\n                  ptr, get(self.registers), atomic, is_smem=False, multimem=True,\n              )\n            else:\n              utils.multimem_store(ptr, get(self.registers))\n          elif atomic is not None:\n            is_smem = utils.is_smem_ref(transfer_ref)","sourceCodeStart":3636,"sourceCodeEnd":3672,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3636-L3672","documentation":"For WGStridedFragLayout, store only supports the 16-byte-swizzle path; any swizzle value other than 16 raises ValueError('Only TiledLayouts support swizzling'). Strided fragment layouts assume 128-bit (16B) aligned vector transactions, and general swizzling modes exist only for TiledLayout.","triggerScenarios":"fa.store(ref, swizzle=32) (or 0/64/128...) on an array with WGStridedFragLayout. swizzle=16 is the only accepted value for this layout; TiledLayout must be used for other swizzle modes.","commonSituations":"Copying TiledLayout store examples with swizzle=32/64 while the array was built with a warpgroup strided layout; TMA store descriptors with non-16B swizzle combined with strided fragments.","solutions":["Use swizzle=16 with WGStridedFragLayout, or pass swizzle=None if no swizzle is needed.","Switch the array to a TiledLayout (fa.to_layout / build it with TiledLayout) when you need 32/64/128-byte swizzle modes."],"exampleFix":"# before\nfa.store(ref, swizzle=32)  # fa has WGStridedFragLayout\n\n# after\nfa.store(ref, swizzle=16)\n# or: fa.to_layout(TiledLayout(...)).store(ref, swizzle=32)","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu.fragmented_array import WGStridedFragLayout\nif isinstance(fa.layout, WGStridedFragLayout):\n    assert swizzle in (16, None), 'WGStridedFragLayout only supports swizzle=16'","typeGuard":"def valid_swizzle(fa, swizzle):\n    if isinstance(fa.layout, WGStridedFragLayout):\n        return swizzle in (16, None)\n    return True","tryCatchPattern":null,"preventionTips":["Use swizzle=16 with strided fragment layouts; use TiledLayout for 32/64/128B swizzle modes.","Centralize swizzle constants next to layout choice instead of hard-coding them at store sites."],"tags":["jax","mosaic-gpu","swizzle","strided-layout","store"],"backgroundTag":"unsupported-layout-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}