[Bug][CUDA] Static shared sub-byte allocation uses floor-sized storage

[LittlehamsterXu]

Summary

The CUDA code generator computes the number of 32-bit storage units for a static shared allocation by dividing the logical element count by 32 / dtype.bits(). The division is floor division, so a non-multiple sub-byte allocation receives zero storage units instead of the one unit required to hold its elements.

On the frozen TVM snapshot, int1[12] and int4[6] lower to __shared__ alignas(64) int s_ptr[0];. The generated kernel then accesses s_ptr[0]. The aligned controls int1[32] and int4[8] lower to [1] and pass the same frontend check.

This is a CUDA backend code-generation observation: the generated declaration has zero elements for inputs whose packed storage requirement is one 32-bit word.

Latest upstream source check

As of 2026-09-05, the same floor division remains in both the upstream main branch and release v0.26.0:

if (is_packed_integer_dtype && scope == "shared") {
  constant_size = constant_size / (32 / dtype.bits());
}

• upstream main
• upstream v0.26.0

This is a source-level version check. The compiler observations above were obtained from the frozen v0.25.0.post1 path; this report does not claim a v0.26.0 binary or device replay.

Environment

• TVM commit: b3e249b7d75f8f3bc7cbee48188d3c80ae323437 (v0.25.0.post1)
• Python: 3.11
• Platform: Ubuntu 22.04 under WSL2, x86_64
• GPU present: NVIDIA GeForce RTX 4070 Laptop GPU, compute capability 8.9

The frozen tvm-eval path used TVM CUDA source generation and Clang 14's nvptx device frontend with a documented minimal qualifier shim. A separate tvm-0.25 environment with apache-tvm 0.25.0.post1 and NVRTC rejects the same unaligned case during tvm.tirx.build() with the size of an array must be greater than zero. Neither path performs device execution.

Affected code

In src/backend/cuda/codegen/codegen_cuda.cc, the frozen generator contains:

size_t constant_size = 1;
for (const auto& dim : op->buffer->shape) {
  const IntImmNode* dim_imm = dim.as<IntImmNode>();
  constant_size *= dim_imm->value;
}
// ...
if ((dtype == DataType::Int(4) || dtype == DataType::UInt(4) || dtype == DataType::Int(1)) &&
    scope == "shared") {
  constant_size = constant_size / (32 / dtype.bits());
}
stream << ' ' << vid << '[' << constant_size << "];\n";

For a one-dimensional shape, the required number of 32-bit storage units is ceil(shape * bits / 32), not floor(shape / (32 / bits)).

Minimal reproduction

Run the following standalone Python program in a TVM build with CUDA code generation. It does not require any project-local file. The program accepts either source-only generation or a compiler rejection caused by the generated zero-length array:

import re

import tvm
from tvm.script import from_source
from tvm.script import ir as I
from tvm.script import tirx as T


ZERO_SIZE_ARRAY = re.compile(
    r"(?:size of an array must be greater than zero|zero[- ]size array|"
    r"zero size arrays)",
    re.IGNORECASE,
)


def module_source(shape, dtype):
    return f'''\
@I.ir_module(s_tir=True)
class Module:
    @T.prim_func(s_tir=True)
    def main(A: T.handle, B: T.handle):
        T.func_attr({{"tirx.noalias": True}})
        a = T.match_buffer(A, [{shape}], "{dtype}", scope="global")
        b = T.match_buffer(B, [{shape}], "{dtype}", scope="global")
        s = T.sblock_alloc_buffer([{shape}], "{dtype}", scope="shared")
        for bx in T.thread_binding(1, thread="blockIdx.x"):
            for tx in T.thread_binding({shape}, thread="threadIdx.x"):
                s[tx] = a[tx]
                b[tx] = s[tx]
'''


reproduced = False
for shape, dtype in [(12, "int1"), (32, "int1"), (6, "int4"), (8, "int4")]:
    module = from_source(
        module_source(shape, dtype),
        extra_vars={"I": I, "T": T},
        check_well_formed=False,
    )
    bits = int(dtype.removeprefix("int"))
    expected_units = (shape * bits + 31) // 32
    try:
        built = tvm.tirx.build(module, target="cuda")
    except Exception as exc:
        detail = str(exc)
        if ZERO_SIZE_ARRAY.search(detail) and "s_ptr[0]" in detail:
            print(
                f"dtype={dtype}, shape={shape}, expected_units={expected_units}, "
                "observed=0, build_status=ZERO_SIZE_ARRAY_REJECTED"
            )
            reproduced = expected_units == 1
            continue
        raise

    source = built.imports[0].inspect_source()
    shared = [line.strip() for line in source.splitlines() if "__shared__" in line]
    match = re.search(r"s_ptr\[(\d+)\];", shared[0]) if shared else None
    observed = int(match.group(1)) if match else None
    print(
        f"dtype={dtype}, shape={shape}, expected_units={expected_units}, "
        f"observed={observed}, build_status=SOURCE_GENERATED"
    )
    print("  " + (shared[0] if shared else "no static shared declaration"))
    if observed is not None and observed < expected_units:
        reproduced = True

if reproduced:
    print("BUG OBSERVED: at least one unaligned sub-byte allocation is undersized")
else:
    print("CANDIDATE NOT OBSERVED")

The reproducer accepts both supported observation modes. In a source-generation-only build it prints the generated s_ptr[0] declaration and continues through all four controls. In an environment where tvm.tirx.build() invokes NVRTC, it catches the compiler rejection, records the zero-size-array diagnostic, and continues with the remaining controls.

Observed from the frozen code generator:

dtype=int1, shape=12, expected_units=1, observed=0
  __shared__ alignas(64) int s_ptr[0];
dtype=int1, shape=32, expected_units=1, observed=1
  __shared__ alignas(64) int s_ptr[1];
dtype=int4, shape=6, expected_units=1, observed=0
  __shared__ alignas(64) int s_ptr[0];
dtype=int4, shape=8, expected_units=1, observed=1
  __shared__ alignas(64) int s_ptr[1];
BUG OBSERVED: at least one unaligned sub-byte allocation is undersized

The generated CUDA body also contains the corresponding accesses:

__shared__ alignas(64) int s_ptr[0];
*(((int*)s_ptr) + ((int)threadIdx.x) / 32) = ...;  // int1[12]
... = *(((int*)s_ptr) + ((int)threadIdx.x) / 32);

For int1[12], every active thread maps to storage word zero; for int4[6], the divisor is eight and every active thread again maps to word zero. Thus the declaration is not dead code or an unused allocation.

Compiler confirmation

The paired controls were fed to Clang 14 with:

clang++ -x cuda --cuda-device-only --cuda-gpu-arch=sm_80 \
  -nocudainc -nocudalib -Werror=zero-length-array -fsyntax-only

Because CUDA headers are unavailable, __global__, __device__, and __shared__ are supplied as a minimal parser shim, and __launch_bounds__ is removed from the isolated kernel input. This preserves the generated array declaration and checks it in Clang's CUDA device frontend:

Case	Generated storage	Reference units	Clang result
int1[12]	[0]	1	rejected: zero-size array
int1[32]	[1]	1	pass
int4[6]	[0]	1	rejected: zero-size array
int4[8]	[1]	1	pass

The observations above were obtained from the generated CUDA source and the compiler frontends described below; the standalone program prints the relevant declaration and diagnostic directly.

With apache-tvm 0.25.0.post1 in the tvm-0.25 environment, the same reproducer also reaches NVRTC. The int1[12] case is rejected with:

tvm_kernels.cu(82): error: the size of an array must be greater than zero
    __shared__ alignas(64) int s_ptr[0];

The updated reproducer treats this compiler rejection as the same source-level undersized-storage observation and preserves the diagnostic in the run log. The NVRTC run still does not establish device execution behavior.

Impact

For affected static CUDA shared allocations, the observed generated declaration may be rejected by the CUDA compiler as a zero-length array. If a compiler accepts zero-length arrays as an extension, the declaration still provides no storage while the kernel performs a word-zero access; the downstream behavior remains to be checked with a supported CUDA toolchain.

Suggested fix

Compute the packed storage count with checked arithmetic and a ceiling operation, for example ceil(element_count * dtype.bits() / 32), or an equivalent integer form that avoids overflow. Add regression tests for unaligned and aligned int1, int4, and uint4 static shared allocations, and compile the generated CUDA source with a supported CUDA toolchain.