This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new b29b52fdf4 [Fix][LLVM] Preserve 64-bit AllocBuffer extents (#20141)
b29b52fdf4 is described below
commit b29b52fdf42e5dd792715e861dee80c6f26c731b
Author: Igor Stadnyk <[email protected]>
AuthorDate: Tue Sep 8 16:54:47 2026 +0100
[Fix][LLVM] Preserve 64-bit AllocBuffer extents (#20141)
Fixes #20125.
## What changed
LLVM CPU code generation now keeps constant `AllocBuffer` extents as
signed 64-bit values through stack allocation emission instead of
narrowing them to `int32_t`. The temporary-allocation alignment helper
accepts the same width and avoids multiplying when a large allocation
cannot reduce the requested alignment, preventing overflow in that
calculation.
This change is intentionally limited to the CPU `CodeGenLLVM` path. The
separate NVPTX and AMDGPU overrides described in the issue remain out of
scope.
## Why
An extent of `2**32 + 1` previously wrapped to one before
`CreateAlloca`, producing a one-element stack allocation followed by an
`i64 4294967296` element access in unoptimized TIR-X LLVM IR. The
generated module could therefore contain an out-of-bounds access even
though the original extent was valid as an `int64_t`.
The regression test only compiles and inspects LLVM IR; it does not
execute or materialize the very large allocation. It covers extents that
previously wrapped to one and four elements.
## Validation
- New regression on the unmodified base: 2 failures
- New regression after the fix: 2 passed
- `tests/python/codegen/test_target_codegen_llvm.py`: 380 passed
- Rebuilt all 65 affected C++ units, including the NVPTX and AMDGPU
callers of the alignment helper
- Changed-file pre-commit hooks and `git diff --check`: passed
- Compile-only boundary probes through `INT64_MAX - 3`: exact `i64`
allocation counts, with no allocation executed
## AI assistance
This PR was prepared with OpenAI Codex. The patch was derived from the
repository and the issue's documented failure mode; no external
third-party code was copied into the change.
Generated-by: OpenAI Codex
Co-authored-by: Igor Stadnyk <[email protected]>
---
src/target/llvm/codegen_llvm.cc | 4 ++--
src/tirx/transform/ir_utils.h | 13 +++++++++----
tests/python/codegen/test_target_codegen_llvm.py | 23 +++++++++++++++++++++++
3 files changed, 34 insertions(+), 6 deletions(-)
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 3ab326c559..2ced6e7847 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2191,7 +2191,7 @@ void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) {
const IntImmNode* dim_imm = op->buffer->shape[0].as<IntImmNode>();
TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack allocation";
- int32_t constant_size = static_cast<int32_t>(dim_imm->value);
+ int64_t constant_size = dim_imm->value;
TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack
allocation";
StorageInfo& info = alloc_storage_info_[op->buffer.get()];
@@ -2206,7 +2206,7 @@ void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) {
info.alignment = 16;
}
llvm::AllocaInst* alloca = WithFunctionEntry([&]() {
- return builder_->CreateAlloca(DTypeToLLVMType(op->buffer->dtype),
ConstInt32(constant_size));
+ return builder_->CreateAlloca(DTypeToLLVMType(op->buffer->dtype),
ConstInt64(constant_size));
});
auto alignment = static_cast<unsigned>(alloca->getAlign().value());
if (alignment < static_cast<unsigned>(info.alignment)) {
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index 9dcd0431a4..9ad1099a91 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -181,12 +181,17 @@ inline PrimType APIType(const PrimType& t) {
* \param const_size The constant size of the array.
* \return the alignment
*/
-inline int GetTempAllocaAlignment(const PrimType& type, int32_t const_size) {
+inline int GetTempAllocaAlignment(const PrimType& type, int64_t const_size) {
int align = runtime::kTempAllocaAlignment;
if (const_size > 0) {
- int64_t const_s = static_cast<int64_t>(const_size) * type.StorageBytes();
- while (align > const_s) {
- align = align / 2;
+ int64_t element_bytes = type.StorageBytes();
+ // Only compute the total size when it can reduce the alignment. This also
avoids
+ // overflowing for very large allocations.
+ if (element_bytes > 0 && const_size <= (align - 1) / element_bytes) {
+ int64_t const_s = const_size * element_bytes;
+ while (align > const_s) {
+ align = align / 2;
+ }
}
}
return align;
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index bdf617ff54..9e8c091310 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -957,6 +957,29 @@ def test_llvm_order_functions():
assert matches == sorted(matches)
[email protected](not env.has_llvm(), reason="need llvm")
[email protected]("extent", [2**32 + 1, 2**32 + 4])
+def test_llvm_large_stack_allocation_uses_64bit_extent(extent):
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((1,), "float32")):
+ B = T.alloc_buffer(
+ (extent,),
+ "float32",
+ scope="global",
+ annotations={"disable_lower_builtin": True},
+ )
+ A[0] = B[extent - 1]
+
+ module = tvm.tirx.build(
+ main,
+ target={"kind": "llvm", "opt-level": 0},
+ pipeline="tirx",
+ )
+ llvm_ir = module.inspect_source("ll")
+
+ assert re.search(rf"alloca float, i64 {extent}(?:,|$)", llvm_ir)
+
+
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@tvm.testing.skip_if_32bit
def test_llvm_import():
