{"record":{"id":"b6ff9cfce0e99931","repo":"jax-ml/jax","slug":"only-byte-aligned-bitcasts-are-supported","errorCode":null,"errorMessage":"Only byte-aligned bitcasts are supported.","messagePattern":"Only byte-aligned bitcasts are supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":1593,"sourceCode":"              f\" dimension, got {ref.shape[0]} != {transformed_shape[0]}.\"\n          )\n        address = arith_dialect.addi(ref.address, _i32_constant(offset))\n        ref = tcgen05.TMEMRef(\n            address=address,\n            shape=cast(tuple[int, int], transformed_shape),\n            dtype=mlir_dtype,\n            layout=layout,\n        )\n      else:\n        assert isinstance(ref, ir.Value)  # make pyrefly happy\n        input_ref_ty = ir.MemRefType(ref.type)\n        if input_ref_ty.memory_space == mgpu_utils.smem():\n          assert layout is None\n          ref_bits = math.prod(transformed_shape) * mgpu_utils.bitwidth(\n              mlir_dtype\n          )\n          if ref_bits % 8:\n            raise NotImplementedError(\"Only byte-aligned bitcasts are supported.\")\n          assert offset % gpu_core.SMEM_ALIGNMENT == 0\n\n          if lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n            if not isinstance(ref.owner, mgpu.dialect.SliceSMEMOp):\n              # This restriction can be lifted by:\n              # - Using memref ops to get the pointer and offset of the base ref.\n              # - Subtracting gpu_dialect.dynamic_shared_memory() from those to\n              #   get the base offset relative to the beginning of SMEM.\n              # - Implementing layout and lowering rules for all ops above.\n              raise NotImplementedError(\n                  \"The base ref for aliases must come from a slice_smem op.\"\n              )\n\n            base_offset = ref.owner.offset.value\n            total_offset = base_offset + offset\n\n            ref_ty = ir.MemRefType.get(\n                transformed_shape, mlir_dtype, memory_space=mgpu_utils.smem()","sourceCodeStart":1575,"sourceCodeEnd":1611,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L1575-L1611","documentation":"When lowering an SMEM alias/bitcast, the total bit count of the transformed region (prod(shape) * bitwidth(dtype)) must be divisible by 8 because memory is byte-addressable. Sub-byte element counts (e.g. i1/i2/i4 blocks) that don't total whole bytes cannot be aliased and raise NotImplementedError.","triggerScenarios":"Aliasing/biting an SMEM ref whose element count times element bitwidth is not a multiple of 8 — e.g. a block of 3 i4 values (12 bits), or boolean arrays of odd length.","commonSituations":"Using sub-byte packed dtypes (i4, i2, u1) in Pallas GPU kernels with view/alias operations; quantization kernels that pack low-bit weights.","solutions":["Pad the block/element count so total bits are byte-aligned","Promote sub-byte dtypes to i8 before aliasing","Repack data so aliased regions start and end on byte boundaries"],"exampleFix":"# before\nb = alloc_smem((3,), jnp.int4))  # 12 bits total\nv = b.view(...)\n# after\nb = alloc_smem((4,), jnp.int4))  # 16 bits, byte-aligned\nv = b.view(...)","handlingStrategy":"validation","validationCode":"total_bits = int(np.prod(shape)) * jnp.dtype(dt).itemsize * 8\nassert total_bits % 8 == 0, f'{total_bits} bits is not byte-aligned'","typeGuard":"def is_byte_aligned(shape, dt) -> bool:\n    return (int(np.prod(shape)) * jnp.dtype(dt).itemsize * 8) % 8 == 0","tryCatchPattern":null,"preventionTips":["Pad sub-byte blocks to whole bytes","Prefer i8 for aliased scratch","Centralize allocation in helpers that enforce alignment"],"tags":["jax","pallas","mosaic-gpu","smem","bitcast","alignment","sub-byte"],"backgroundTag":"memory-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}