ruiling-smartbear opened a new pull request, #20384: URL: https://github.com/apache/tvm/pull/20384
Fixes #20269. Since the BigInt migration the element count of a constant shape is formed exactly, but the sites that turn it into a bit or byte count still used unchecked 64-bit arithmetic, so a shape whose size does not fit wrapped around and was then used as if it were a real size. Add a checked helper to `tirx/transform/ir_utils.h` that forms the whole product exactly and reports when it does not fit in `uint64_t`, and route the reported sites through it. - `LowerTVMBuiltin` folds the byte count into the `uint64` argument of `TVMBackendAllocWorkspace`; constant folding wraps it modulo `2**64`, so the shapes in the report request 0 bytes, or `2**62` instead of `5 * 2**62`. The stack-allocation check multiplied the element count by the element size in `size_t` in the same way. A constant size that does not fit is now rejected instead of emitted wrapped. A shape with symbolic extents still builds the multiplication chain and is unchanged; guarding that needs a runtime check in the generated code. - `StoragePlanRewriter::FindAlloc` multiplied the element count by the element bits in `uint64`. A wrapped product can look like a small known size, so a buffer is placed in an allocation far too small for it: `8 * (2**61 + 128)` wraps to 1024 bits, and the planner put that buffer inside a 128-byte one. A constant size that does not fit is now rejected, the same way `LowerTVMBuiltin` rejects it. Leaving it unplanned instead would make the tagged-memory path fail later on the unrelated `Special tagged memory must be const size` check. The `const_nbits * match_range` end of the free-list search is saturated: wrapping put the upper bound below the lower one, and the search then ran past the end of `const_free_map_`. - The inplace check in the same pass compared the same wrapped product against the source entry, and no longer matches when the size does not fit. - `VerifyGPUCode` converted the exact element count to `int64_t`, which now throws `OverflowError` from inside the analysis instead of reporting a violation, and then multiplied by the element bytes in `size_t`: a shared buffer of `2**31 x 2**31` float32 wraps to zero, and the kernel is accepted. The size is now reported as an error, and the per-block totals accumulate without wrapping. - `MergeSharedMemoryAllocations` had its own copy of the element-count product followed by the same unchecked multiply. It now uses the shared helper and rejects a size that does not fit, and its search range is saturated the same way. `AllocBuffer::ConstantAllocationSize()` is not touched: it already reports an element count that does not fit, but the callers need the size in bits or bytes, and that is where the remaining overflow was. ## Validation Linux x86_64, CPU-only build with LLVM 18 (no CUDA), on `850a9d8`: - New tests on the unpatched tree: `test_workspace_size_that_does_not_fit_is_rejected` (3 shapes) and `test_allocation_size_that_does_not_fit_is_rejected` fail with the wrapped size accepted, `test_shared_memory_size_that_does_not_fit[int8]` fails with `BigInt does not fit int64_t` raised from the analysis, and the float32 case is accepted; this is the report reproduced. The within-limit cases pass. With the fix all 10 pass. `test_reuse_search_range_does_not_wrap` also passes before the fix: with libstdc++, incrementing `end()` happens to land on the only free entry, so it guards the invariant rather than reproducing a crash. - `tests/python/tirx-transform/`: 340 passed, 8 xfailed, 1 xpassed. `tests/python/s_tir/transform/`: 268 passed, 8 skipped. `tests/python/s_tir/analysis/`: 92 passed. `tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py`: 8 passed. - Changed-file pre-commit checks passed. No GPU tests were run. -- 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]
