[Bug] TIR allocation-size analysis silently overflows in `ConstantAllocationSize` and `CalculateAllocatedBytes`

[LittlehamsterXu]

Summary

TVM v0.25.0.post1 performs unchecked signed int64_t arithmetic when calculating the size of constant AllocBuffer shapes. When the product exceeds INT64_MAX, the result wraps (or invokes signed-overflow undefined behavior) instead of being reported as unknown or rejected.

This is separate from apache/tvm#20125: #20125 is an LLVM codegen int64_t-to-int32_t narrowing bug, while this issue affects TIR analysis and allocation planning before codegen.

Environment

• TVM source: local clone of Apache TVM
• TVM source tag: v0.25.0.post1
• TVM commit: b3e249b7d75f8f3bc7cbee48188d3c80ae323437 (v0.25.0.post1)
• Python: 3.11.15
• TVM package: apache-tvm 0.25.0.post1
• NumPy: 2.4.6
• pytest: 8.4.2
• 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

The test was run in the Conda environment tvm-0.25.

Affected code

• include/tvm/tirx/stmt.h L292-L301
  • AllocBuffer::ConstantAllocationSize() multiplies shape extents in int64_t without overflow checking.
• src/s_tir/analysis/calculate_allocated_memory.cc L70-L81
  • AllocBufferCalculator::VisitStmt_() repeats the unchecked shape multiplication and then unchecked byte-size multiplication/addition.
• src/s_tir/analysis/calculate_allocated_memory.cc L103-L121
  • The two CalculateAllocatedBytes() overloads expose the helper's wrapped result.

The return type of ConstantAllocationSize() is std::optional<int64_t>, but arithmetic overflow does not produce std::nullopt; it produces a wrapped value. In CalculateAllocatedBytes(), a negative wrapped value can also be hidden because _max_size is initialized to 0 and updated with std::max(_current_size, _max_size).

Minimal reproduction

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

cases = [
    ((2**62, 4), "int8"),
    ((2**62, 5), "int8"),
    ((2**32, 2**32), "int8"),
]

for shape, dtype in cases:
    buf = tvm.tirx.decl_buffer(shape, dtype=dtype, scope="global.vtcm")
    func = tvm.tirx.PrimFunc([], tvm.tirx.AllocBuffer(buf))
    got = tvm.s_tir.analysis.calculate_allocated_bytes(func)["main"]["global.vtcm"]
    expected = shape[0] * shape[1]
    print(f"shape={shape}, expected={expected}, got={got}")
PY

Observed output:

shape=(4611686018427387904, 4), expected=18446744073709551616, got=0
shape=(4611686018427387904, 5), expected=23058430092136939520, got=4611686018427387904
shape=(4294967296, 4294967296), expected=18446744073709551616, got=0

The expected values are not representable in signed int64_t; the API currently returns wrapped values instead of rejecting the allocation or returning an unknown result.

Impact

• CalculateAllocatedBytes() can under-report memory usage.
• s_tir.transform.VerifyVTCMLimit consumes this result, so a wrapped value can bypass or weaken a VTCM limit check.
• Storage-planning passes that consume ConstantAllocationSize() can classify an oversized allocation as a valid, differently sized constant allocation.
• The behavior is undefined at the C++ level because signed overflow is not defined by the language standard.

The existing allocation-analysis test passes (6 passed), but it does not cover overflow boundaries.

Suggested fix

1. Use checked multiplication and addition for shape elements, dtype bytes, and accumulated scope sizes.
2. Reject negative extents at the allocation boundary.
3. Define one consistent overflow result (std::nullopt, a diagnostic, or an explicit saturated/unknown value) and preserve it through all callers.
4. Add boundary tests for products just below/at/above INT64_MAX and for wrapped-to-zero products.