{"record":{"id":"85b8a72740507490","repo":"jax-ml/jax","slug":"async-store-smem-requires-a-tiled-and-swizzled-ref","errorCode":null,"errorMessage":"async_store_smem requires a tiled and swizzled ref","messagePattern":"async_store_smem requires a tiled and swizzled ref","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":676,"sourceCode":"    atomic_type = None\n    if atomic is not None:\n      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))\n    mgpu.dialect.async_store_smem(\n        src,\n        ref_smem,\n        barrier.as_barrier_memref(),\n        gpu_cluster_dim.value,\n        cluster_idx_i32,\n        atomic_type=atomic_type,\n        optimized=optimized,\n    )\n    return ()\n\n  match remaining_ref_transforms:\n    case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntilingTransform(tiling)):\n      pass\n    case _:\n      raise NotImplementedError(\"async_store_smem requires a tiled and swizzled ref\")\n\n  total_bits = math.prod(shape) * dtypes.itemsize_bits(dtype)\n  if total_bits % 8:\n    raise ValueError(\n        f\"Can only transfer integer bytes (shape={shape}, dtype={dtype})\"\n    )\n  total_bytes = total_bits // 8\n  if total_bytes % WARPGROUP_SIZE:\n    raise NotImplementedError(f\"Transfer is not a multiple of {WARPGROUP_SIZE} bytes\")\n\n  peer_barrier = barrier.remap_to_cluster(gpu_cluster_dim, cluster_idx_val)\n  peer_barrier.arrive_expect_tx(total_bytes // WARPGROUP_SIZE)\n\n  lowering._ensure_fa(src, dtype).store_tiled_async(\n      ref_smem,\n      barrier,\n      cluster_dim=gpu_cluster_dim,\n      cluster_idx=cluster_idx_val,","sourceCodeStart":658,"sourceCodeEnd":694,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L658-L694","documentation":"In the Warp lowering path, a match statement requires the remaining ref transforms to be exactly (UnswizzleRef(swizzle), UntilingTransform(tiling)) — i.e. the target must be a tiled and swizzled SMEM ref. Any other combination raises NotImplementedError because the async TMA store instruction only works with that layout.","triggerScenarios":"Calling async_store_smem on an SMEM ref that is not both tiled (untiling transform present) and swizzled, e.g. a plain contiguous SMEM block, or one that is tiled but not swizzled.","commonSituations":"Writing a custom pipeline (e.g. GEMM epilogues, attention) and forgetting to construct the SMEM ref via the tiled/swizzled helpers; using MMA output layouts that skip the swizzle step.","solutions":["Construct the destination with the standard tiling+swizzle transforms expected by the pipeline (e.g. use the layouts produced by mgpu/pl.tiled helpers)","Copy to a properly tiled+swizzled intermediate SMEM buffer and store from there","Fall back to a synchronous store if layout flexibility is more important than async overlap"],"exampleFix":"# before\nplain_smem = pl.SMEM((128, 128), dtype)\nasync_store_smem(plain_smem, x, barrier)\n# after\ntiled_smem = make_tiled_swizzled_smem((128, 128), dtype)  # tiled + swizzled\nasync_store_smem(tiled_smem, x, barrier)","handlingStrategy":"fallback","validationCode":"from jax._src.pallas import gpu_core\ndef is_tiled_and_swizzled(transforms):\n  return (len(transforms) == 2 and\n          isinstance(transforms[0], gpu_core.UnswizzleRef) and\n          isinstance(transforms[1], gpu_core.UntilingTransform))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build destination SMEM refs through the tiling/swizzle helpers used by the examples","Fall back to synchronous stores when the layout can't be made tiled+swizzled"],"tags":["jax","pallas","mosaic-gpu","tiled-layout","swizzle","async-store"],"backgroundTag":"unsupported-memory-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}