{"record":{"id":"768ff1361f5e09ee","repo":"jax-ml/jax","slug":"unsupported-swizzle-expected-16-32-64-or-128-b","errorCode":null,"errorMessage":"Unsupported swizzle, expected 16, 32, 64 or 128, but got: {swizzle}","messagePattern":"Unsupported swizzle, expected 16, 32, 64 or 128, but got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":2006,"sourceCode":"    nvvm.tcgen05_cp(\n        nvvm.Tcgen05CpShape.SHAPE_128x128b, ptr, desc,\n        group=nvvm.CTAGroupKind.CTA_2 if collective else nvvm.CTAGroupKind.CTA_1\n    )\n\n\ndef async_copy_smem_to_tmem(\n    smem_ref: ir.Value,\n    tmem_ref: TMEMRef,\n    swizzle: int,\n    collective: bool = False,\n) -> None:\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  smem_ty = ir.MemRefType(smem_ref.type)\n  if (dtype := smem_ty.element_type) != tmem_ref.dtype:\n    raise ValueError(f\"Incompatible dtypes: SMEM has {dtype}, TMEM has {tmem_ref.dtype}\")\n  if swizzle not in {16, 32, 64, 128}:\n    raise ValueError(f\"Unsupported swizzle, expected 16, 32, 64 or 128, but got: {swizzle}\")\n  bitwidth = utils.bitwidth(dtype)\n  if tmem_ref.packing != 32 // bitwidth:\n    raise ValueError(\n        \"tcgen05.cp only supports fully packed TMEM references\"\n        f\" (packing={32 // bitwidth}), but got packing={tmem_ref.packing}\"\n    )\n  if tmem_ref.shape[0] != TMEM_ROWS:\n    raise ValueError(\n        f\"TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref.shape[0]}\"\n    )\n  if tmem_ref.layout != tmem_default_layout(packing=tmem_ref.packing):\n    raise ValueError(\n        f\"Only standard TMEM layout is supported, got: {tmem_ref.layout}\"\n    )\n  swizzle_elems = 8 * swizzle // bitwidth\n  expected_smem_shape = utils.tile_shape(tmem_ref.shape, (8, swizzle_elems))\n  smem_shape = tuple(smem_ty.shape)\n  if smem_shape != expected_smem_shape:","sourceCodeStart":1988,"sourceCodeEnd":2024,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1988-L2024","documentation":"Raised by async_copy_smem_to_tmem when the swizzle parameter is not one of the hardware-supported swizzle sizes. NVIDIA tcgen05.cp only implements 16-, 32-, 64- and 128-byte swizzle modes for the SMEM-to-TMEM path; anything else (including 0 for no swizzle) cannot be lowered.","triggerScenarios":"Calling tcgen05.async_copy_smem_to_tmem(..., swizzle=N) with N not in {16, 32, 64, 128} — e.g. swizzle=0 for unswizzled access, or swizzle=256 copied from a TMA descriptor config.","commonSituations":"Porting kernels that previously used unswizzled SMEM loads, or reusing TMA descriptor swizzle values (which also only allow 32/64/128) incorrectly; misreading swizzle as bits vs bytes.","solutions":["Use one of the supported swizzle values: 16, 32, 64, or 128","Match the swizzle to the SMEM layout you allocated (the SMEM shape check later enforces swizzle_elems = 8*swizzle/bitwidth consistency)","If you truly need unswizzled copies, use a different copy path (e.g. regular async_copy) instead of tcgen05.cp"],"exampleFix":"// before\nasync_copy_smem_to_tmem(smem, tmem, swizzle=0)\n// after\nasync_copy_smem_to_tmem(smem, tmem, swizzle=128)","handlingStrategy":"validation","validationCode":"assert swizzle in {16, 32, 64, 128}, f'swizzle {swizzle} unsupported for tcgen05.cp'","typeGuard":"def is_valid_swizzle(s):\n    return s in {16, 32, 64, 128}","tryCatchPattern":null,"preventionTips":["Define the swizzle enum once per kernel and reuse it for both SMEM layout and copy calls","Remember swizzle is in bytes; never pass 0 to tcgen05.cp paths"],"tags":["gpu","mosaic","tcgen05","swizzle","validation"],"backgroundTag":"invalid-parameter-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}