LittlehamsterXu opened a new issue, #20268: URL: https://github.com/apache/tvm/issues/20268
# [Bug] TIR allocation-size analysis silently overflows in `ConstantAllocationSize` and `CalculateAllocatedBytes` ## 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](https://github.com/apache/tvm/issues/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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/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 ```bash 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: ```text 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
