[Bug] Allocation/shape arithmetic wraps in storage planning and workspace lowering
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
## 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](https://github.com/apache/tvm/issues/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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/tirx/transform/storage_rewrite.cc#L894-L899), [`L973-L980`](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/tirx/transform/storage_rewrite.cc#L973-L980), and [`L1001-L1005`](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/tirx/transform/storage_rewrite.cc#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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/tirx/transform/lower_tvm_builtin.cc#L263-L265), [`L286-L289`](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/tirx/transform/lower_tvm_builtin.cc#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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/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](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/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`:
```bash
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:
```text
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.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by running the provided Python reproduction for LowerTVMBuiltin, then inspect the cited paths in storage_rewrite.cc, lower_tvm_builtin.cc, verify_gpu_code.cc, and merge_shared_memory_allocations.cc. Compare their allocation-size arithmetic and add regression coverage for products wrapping to zero or a smaller positive value. Done means affected paths consistently detect or propagate overflow rather than under-counting allocation sizes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100