{"record":{"id":"50795f2347519726","repo":"jax-ml/jax","slug":"the-smem-tiles-must-be-contiguous","errorCode":null,"errorMessage":"The SMEM tiles must be contiguous","messagePattern":"The SMEM tiles must be contiguous","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":2032,"sourceCode":"    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:\n    raise ValueError(\n        f\"SMEM has shape {smem_shape}, but expected {expected_smem_shape} for\"\n        f\" TMEM shape {tmem_ref.shape} with swizzle={swizzle}\"\n    )\n  strides, _ = smem_ty.get_strides_and_offset()\n  row_tile_stride, col_tile_stride, inner_row_stride, inner_col_stride = strides\n  if inner_col_stride != 1 or inner_row_stride != swizzle_elems:\n    raise ValueError(\"The SMEM tiles must be contiguous\")\n  # Make sure strides are a multiple of the byte packing for narrow types.\n  byte_packing = max(8 // bitwidth, 1)\n  assert row_tile_stride % byte_packing == 0\n  assert col_tile_stride % byte_packing == 0\n\n  # Figure out the matrix descriptor parameters (LBO/SBO)\n  # The copy happens using the usual \"core matrix\" structure: a memory region\n  # describing a 8x128bit matrix. LBO describes how far apart from each other\n  # are consecutive matrices along the minor dimension (in our case the minor\n  # dim is contiguous, so exactly 128 bit = 16 bytes apart). SBO describes how\n  # far apart is the beginning of the next matrix along the major dimension.\n  # We use a tiling of 8, so it is simply the tile stride.\n  leading_byte_offset = 16\n  stride_byte_offset = row_tile_stride * bitwidth // 8\n  assert tmem_ref.shape[1] * bitwidth // 8 >= 16\n  if swizzle == 16:\n    cp_shape = nvvm.Tcgen05CpShape.SHAPE_128x128b\n    cp_cols_bytes = 16  # 128 bit = 16 bytes","sourceCodeStart":2014,"sourceCodeEnd":2050,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L2014-L2050","documentation":"Raised by async_copy_smem_to_tmem when, after the shape check passes, the SMEM memref strides are not the contiguous strides of the tiled shape: the innermost (column) stride must be 1 and the inner row stride must equal swizzle_elems. This catches views (slices/transposes) of a correctly-shaped buffer that broke contiguity of the 8x(swizzle_elems) core matrices.","triggerScenarios":"Passing a transposed, sliced, or padded SMEM memref to async_copy_smem_to_tmem — e.g. smem[:, ::2] or a subview with inner_col_stride != 1, or an inner row stride that no longer equals 8*swizzle/bitwidth.","commonSituations":"Reusing one big SMEM buffer with strided sub-allocations per pipeline stage; creating the SMEM ref via memref.subview with offsets/strides for double buffering.","solutions":["Allocate a dedicated contiguous SMEM buffer with the exact tiled shape instead of a strided subview","If subview is needed for double buffering, make the tile stride change only the outer (row_tile/col_tile) strides, keeping inner strides (1, swizzle_elems) intact","Copy/reshape the data into a contiguous buffer before the tcgen05.cp"],"exampleFix":"# before\nsmem_view = memref.subview(big_smem, offsets, strides=[1, 2, 2, 1], ...)\n# after\nstage_smem = smem_alloc(f32, smem_shape)  # contiguous, own buffer per stage\n# store into stage_smem, then copy from it","handlingStrategy":"validation","validationCode":"bw = utils.bitwidth(dtype)\nswizzle_elems = 8 * swizzle // bw\nstrides, _ = ir.MemRefType(smem_ref.type).get_strides_and_offset()\nassert strides[3] == 1 and strides[2] == swizzle_elems, 'SMEM tiles not contiguous'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid strided subviews of SMEM for tcgen05.cp; use fresh contiguous allocations per pipeline stage","If double buffering, scale only the outer tile strides and keep inner strides at (swizzle_elems, 1)"],"tags":["gpu","mosaic","tcgen05","strides","shared-memory","contiguity"],"backgroundTag":"non-contiguous-memory-view","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}