[Fix][LLVM] Preserve 64-bit AllocBuffer extents

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)) {

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;

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():