{"record":{"id":"49c24c362bc21ff2","repo":"jax-ml/jax","slug":"transfer-of-total-bits-bits-is-not-divisible-by","errorCode":null,"errorMessage":"Transfer of {total_bits} bits is not divisible by {8 * utils.WARPGROUP_SIZE}","messagePattern":"Transfer of (.+?) bits is not divisible by (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":619,"sourceCode":"    raise NotImplementedError(f\"Expected TiledLayout, got {type(layout)}\")\n\n  ref = op.destination\n  transforms_attr = inference_utils.in_transforms(op)[0]\n  swizzle = swizzle_from_transforms_attr(transforms_attr)\n  unwrapped_ref = unwrap_transformed_memref(ref, transforms_attr)\n  tiling_transform, = memref_transforms_from_transforms_attr(transforms_attr)\n  assert isinstance(tiling_transform, lc.TileTransform)\n\n  dialect_barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)\n  barrier_ref = dialect_barrier.barrier_ref\n\n  cluster_dim = gpu.Dimension(op.cluster_dim.value)  # pyrefly: ignore[missing-attribute]\n  cluster_idx = arith.index_cast(index, op.cluster_idx)\n  cluster_barrier_ref = barrier_ref.remap_to_cluster(cluster_dim, cluster_idx)\n\n  total_bits = math.prod(value.shape) * utils.bitwidth(value.mlir_dtype)\n  if total_bits % (8 * utils.WARPGROUP_SIZE):\n    raise NotImplementedError(\n        f\"Transfer of {total_bits} bits is not divisible by \"\n        f\"{8 * utils.WARPGROUP_SIZE}\"\n    )\n  cluster_barrier_ref.arrive_expect_tx(total_bits // 8 // utils.WARPGROUP_SIZE)\n\n  atomic = None\n  if op.atomic_type is not None:\n    atomic = str(mgpu.AtomicOpType(op.atomic_type.value))  # pyrefly: ignore[missing-attribute]\n\n  def store_tiled_async(optimized: bool):\n    value.store_tiled_async(\n        unwrapped_ref,\n        barrier_ref,\n        cluster_dim=cluster_dim,\n        cluster_idx=cluster_idx,\n        swizzle=swizzle.value if swizzle != mgpu.SwizzlingMode.kNoSwizzle else None,\n        optimized=optimized,\n        tiling_rank=len(tiling_transform.tiling),","sourceCodeStart":601,"sourceCodeEnd":637,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L601-L637","documentation":"async_store_smem uses arrive_expect_tx to have the cluster barrier count incoming bytes per warp; the total transfer size in bits must be divisible by 8*WARPGROUP_SIZE (i.e. bytes divisible by 128 per warpgroup) or the TMA transaction count cannot be computed.","triggerScenarios":"Storing a tensor whose total elements * dtype bitwidth is not divisible by 8 * WARPGROUP_SIZE (1024 bits on 128-thread warpgroups).","commonSituations":"Small or odd-shaped stores (e.g. f16 tensors with a non-multiple-of-64 element count) to SMEM via async_store_smem with cluster barriers enabled.","solutions":["Pad the stored value's shape so total bits are divisible by 8*WARPGROUP_SIZE","Use vector.store / regular store to SMEM instead of async_store_smem for odd sizes","Choose a dtype/tile size whose byte size is a multiple of 128 per warpgroup"],"exampleFix":"// before\nasync_store_smem(x, ref)  # x has 30 f32 elements\n// after\nx = pad_to_multiple(x, 32)  # total bytes divisible by 128\nasync_store_smem(x, ref)","handlingStrategy":"validation","validationCode":"import math\nfrom jax.experimental.mosaic.gpu import utils\ntotal_bits = math.prod(value.shape) * utils.bitwidth(dtype)\nassert total_bits % (8 * utils.WARPGROUP_SIZE) == 0, 'pad store to warpgroup multiple'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad stored tiles so total bytes are divisible by 128 per warpgroup","Prefer power-of-two tile shapes"],"tags":["jax","mosaic","gpu","smem","alignment","tma"],"backgroundTag":"size-alignment-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}