Skip to content

[Fix][TIR] Check allocation size arithmetic instead of wrapping around - #20384

Closed
ruiling-smartbear wants to merge 1 commit into
apache:mainfrom
ruiling-smartbear:fix/20269-checked-allocation-size
Closed

ruiling-smartbear wants to merge 1 commit into
apache:mainfrom
ruiling-smartbear:fix/20269-checked-allocation-size

Conversation

@ruiling-smartbear

Copy link
Copy Markdown

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.

The element count of a constant shape is now formed exactly, but the
sites that multiply it by the element size still used unchecked 64-bit
arithmetic, so a shape whose size does not fit wrapped around and was
then used as a real size.

Add a checked helper to tirx/transform/ir_utils.h and route the reported
sites through it: the workspace byte count and stack-allocation check in
LowerTVMBuiltin, the constant bit count and free-list search range in
StoragePlanRewriter, the same product in MergeSharedMemoryAllocations,
and the shape product and memory totals in VerifyGPUCode.
@tqchen tqchen closed this Sep 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants