LittlehamsterXu opened a new issue, #20269: URL: https://github.com/apache/tvm/issues/20269
## 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. -- 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]
