[Bug] Allocation/shape arithmetic wraps in storage planning and workspace lowering

[LittlehamsterXu]

Summary

Several TVM v0.25.0.post1 allocation paths propagate unchecked or already-wrapped shape sizes into storage planning, GPU verification, and workspace allocation. This is related to apache/tvm#20125 and to the separate ConstantAllocationSize/CalculateAllocatedBytes analysis issue, but covers downstream arithmetic and unsigned propagation rather than LLVM's int64_t-to-int32_t narrowing.

Environment

• TVM source: local clone of Apache TVM
• TVM source tag: v0.25.0.post1
• TVM commit: b3e249b7d75f8f3bc7cbee48188d3c80ae323437 (v0.25.0.post1)
• Conda environment: tvm-0.25
• Python: 3.11.15
• TVM package: apache-tvm 0.25.0.post1
• NumPy: 2.4.6
• Platform: Ubuntu 22.04 under WSL2, x86_64
• Enabled TVM targets: llvm; cuda; nvptx
• GPU present: NVIDIA GeForce RTX 4070 Laptop GPU, driver 591.74

Affected code

• src/tirx/transform/storage_rewrite.cc L894-L899, L973-L980, and L1001-L1005
  • Converts ConstantAllocationSize() to uint64_t, multiplies by element bits, and multiplies by match_range without checking overflow.
• src/tirx/transform/lower_tvm_builtin.cc L263-L265, L286-L289
  • Builds total_bytes as UInt(64) and repeatedly multiplies by shape extents. The resulting UInt64 expression can wrap before being passed to TVMBackendAllocWorkspace.
• src/s_tir/analysis/verify_gpu_code.cc L70-L84
  • Repeats unchecked signed shape multiplication and then casts the result to size_t for local/shared-memory accounting.
• src/s_tir/transform/merge_shared_memory_allocations.cc L67-L72
  • Performs signed multiplication before the result > INT64_MAX check. Once signed overflow has wrapped the value into the int64_t range, this check is always false and cannot detect the overflow.

Reproduction: wrapped workspace byte count

The following applies LowerTVMBuiltin directly to a CPU-targeted AllocBuffer:

conda activate tvm-0.25
python - <<'PY'
import tvm

for shape in [(2**32, 2**32), (2**62, 4), (2**62, 5)]:
    buf = tvm.tirx.decl_buffer(shape, dtype="int8", scope="global")
    alloc = tvm.tirx.AllocBuffer(buf)
    body = tvm.tirx.AttrStmt(
        tvm.tirx.Var("dev", "int32"),
        "device_id",
        tvm.tirx.IntImm("int32", 0),
        alloc,
    )
    attrs = tvm.ir.DictAttrs({"target": tvm.target.Target("llvm")})
    func = tvm.tirx.PrimFunc([], body, attrs=attrs)
    lowered = tvm.tirx.transform.LowerTVMBuiltin()(tvm.IRModule({"main": func}))
    print("shape", shape)
    print(lowered["main"])
PY

Relevant observed output:

shape (4294967296, 4294967296)
T.TVMBackendAllocWorkspace(1, 0, T.uint64(0), 0, 8)

shape (4611686018427387904, 4)
T.TVMBackendAllocWorkspace(1, 0, T.uint64(0), 0, 8)

shape (4611686018427387904, 5)
T.TVMBackendAllocWorkspace(1, 0, T.uint64(4611686018427387904), 0, 8)

For dtype=int8, the first two shapes require 2**64 bytes/elements before overflow, but the generated workspace request is 0. The third requires 5 * 2**62, but the generated request is only 2**62 after wraparound.

Impact

• Storage reuse and inplace-allocation matching can use wrapped const_nbits values.
• GPU local/shared-memory verification can under-count usage after signed-to-unsigned conversion.
• Workspace allocation can receive a zero or undersized byte count for a shape whose mathematical size is much larger.
• If the resulting buffer is accessed, this can become an allocation failure, wrong result, or out-of-bounds access depending on the pipeline and runtime allocator behavior.

This issue is broader than #20125: #20125 is a codegen-width truncation; this issue is arithmetic overflow/wrap across planning and lowering. The two problems can overlap on the same malformed or extreme input.

Suggested fix

1. Introduce a shared checked allocation-size helper for shape-product, dtype-size, and accumulation arithmetic.
2. Check for non-negative extents and overflow before constructing UInt(64) byte expressions.
3. Do not cast a possibly negative signed result to uint64_t/size_t.
4. Make storage planning, GPU verification, and workspace lowering consistently reject or propagate an unknown/overflowed allocation size.
5. Add tests covering products that wrap to zero and products that wrap to a smaller positive value.