apache / apache/tvm

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

Open
#20,269 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.