[Fix][TIR] Check allocation size arithmetic instead of wrapping around - #20384
Closed
ruiling-smartbear wants to merge 1 commit into
Closed
ruiling-smartbear wants to merge 1 commit into
ruiling-smartbear wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.hthat forms the whole product exactly and reports when it does not fit inuint64_t, and route the reported sites through it.LowerTVMBuiltinfolds the byte count into theuint64argument ofTVMBackendAllocWorkspace; constant folding wraps it modulo2**64, so the shapes in the report request 0 bytes, or2**62instead of5 * 2**62. The stack-allocation check multiplied the element count by the element size insize_tin 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::FindAllocmultiplied the element count by the element bits inuint64. 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 wayLowerTVMBuiltinrejects it. Leaving it unplanned instead would make the tagged-memory path fail later on the unrelatedSpecial tagged memory must be const sizecheck. Theconst_nbits * match_rangeend of the free-list search is saturated: wrapping put the upper bound below the lower one, and the search then ran past the end ofconst_free_map_.VerifyGPUCodeconverted the exact element count toint64_t, which now throwsOverflowErrorfrom inside the analysis instead of reporting a violation, and then multiplied by the element bytes insize_t: a shared buffer of2**31 x 2**31float32 wraps to zero, and the kernel is accepted. The size is now reported as an error, and the per-block totals accumulate without wrapping.MergeSharedMemoryAllocationshad 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:test_workspace_size_that_does_not_fit_is_rejected(3 shapes) andtest_allocation_size_that_does_not_fit_is_rejectedfail with the wrapped size accepted,test_shared_memory_size_that_does_not_fit[int8]fails withBigInt does not fit int64_traised 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_wrapalso passes before the fix: with libstdc++, incrementingend()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.