From 4c0b318ef74696772e6353d28b1c1f02c6b66b95 Mon Sep 17 00:00:00 2001
From: Igor Stadnyk <fallenmi@users.noreply.github.com>
Date: Sun, 16 Aug 2026 07:57:37 +0100
Subject: [PATCH] [Fix][LLVM] Preserve 64-bit AllocBuffer extents

LLVM code generation narrowed constant AllocBuffer extents to int32 before creating stack allocations. Preserve the signed 64-bit extent through alignment calculation and LLVM alloca emission, while avoiding overflow in the alignment helper.

Add compile-only regression coverage for large extents that previously wrapped to one or four elements.

Generated-by: OpenAI Codex
---
 src/target/llvm/codegen_llvm.cc               |  4 ++--
 src/tirx/transform/ir_utils.h                 | 13 +++++++----
 .../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 dc454de5ad3c..601bbb32ad47 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2118,7 +2118,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()];
@@ -2133,7 +2133,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 04da364c90a4..b01df7d606e2 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -180,12 +180,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 6f9f128a9828..104b8c520522 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -957,6 +957,29 @@ def Kirby(v: T.float32) -> T.float32:
     assert matches == sorted(matches)
 
 
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+@pytest.mark.parametrize("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():