[Codegen][LLVMCPU] Hoist in-loop stack alloca (#24749)
`HALDispatchABI::loadProcessorData` built its CPU-feature patch buffer
with an `alloca` at the ukernel call site instead of the function entry
block. `RewriteCallOpABI` materializes this after all MLIR->LLVM
conversions are done, so no later pass can hoist it out of a loop if the
call happens to sit inside one, causing the stack to grow unboundedly on
large matmuls.
The computation only depends on the function's `environment` argument
and static target attributes, so it can safely move to the entry block.
Fixes #24744
---------
Signed-off-by: Federico Bruzzone <federico.bruzzone.i@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp
index 6e49965..4c04a5d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/DispatchABI.cpp
@@ -992,6 +992,20 @@
// To get a pointer to the processor data we need to track pointers all the
// way from the environment argument. This is redundant with loadFieldValue
// but that returns values instead.
+ //
+ // `forOp` here is the call itself, which may sit inside a loop, so if we
+ // built these ops at its insertion point we'd re-run the stack allocation
+ // on every loop iteration with nothing to ever pop it back off, overflowing
+ // the stack for large iteration counts (#24744). Build them in the
+ // enclosing function's entry block instead, where they only run once.
+ // The computation only depends on the function's `environment`
+ // argument and static target attributes, so it can safely move to the
+ // enclosing function's entry block instead, where it only runs once.
+ auto funcOp = forOp->getParentOfType<LLVM::LLVMFuncOp>();
+ assert(funcOp && "usage requires an enclosing LLVMFuncOp");
+ OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToStart(&funcOp.getFunctionBody().front());
+
auto loc = forOp->getLoc();
auto environmentPtrValue =
buildArgDI(forOp, /*argNum=*/0, getLocalArgument(forOp, 0), "environment",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir
index 1049e32..e3bfe93 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir
@@ -43,6 +43,59 @@
// -----
+// Regression test for #24744: a `hal.import.bitcode` call requesting
+// "processor_data" (e.g. a ukernel call) gets a local `i64` buffer allocated
+// to patch in statically-known CPU-feature bits. If that alloca were built at the call
+// site instead of the function entry block, it would re-execute on every
+// loop iteration without ever popping the stack back off, overflowing it for
+// large iteration counts. Checks each of the 4 blocks below: the alloca must
+// be in the entry block and must not reappear in the loop body.
+#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+dotprod", target_triple = "aarch64-none-elf"}>
+module {
+ func.func private @default_cconv_with_extra_fields_in_loop(memref<f32>, i32, f64) -> (f32) attributes {
+ hal.import.bitcode = true,
+ hal.import.cconv = 0 : i32,
+ hal.import.fields = ["processor_data", "processor_id"],
+ llvm.bareptr = true
+ }
+ func.func @loop_caller() attributes {hal.executable.target = #executable_target} {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 1024 : index
+ %step = arith.constant 1 : index
+ %c0 = arith.constant 42 : i32
+ %c1 = arith.constant 42.0 : f64
+ %0 = memref.alloca() : memref<f32>
+ scf.for %i = %lb to %ub step %step {
+ %1 = func.call @default_cconv_with_extra_fields_in_loop(%0, %c0, %c1) : (memref<f32>, i32, f64) -> (f32)
+ }
+ return
+ }
+}
+// CHECK: llvm.func @loop_caller
+// Entry block: the processor-data patch buffer is built once, here.
+// CHECK-NOT: ^{{.+}}:
+// CHECK: %[[ENV_DATA:.+]] = llvm.getelementptr inbounds %arg0[{{[0-9]+}}]
+// CHECK-NOT: ^{{.+}}:
+// CHECK: %[[PATCHED_DATA:.+]] = llvm.alloca %{{.+}} x i64
+// CHECK-NOT: ^{{.+}}:
+// CHECK: llvm.load %[[ENV_DATA]]
+// CHECK-NOT: ^{{.+}}:
+// CHECK: llvm.br ^[[HEADER:.+]](
+// Loop header: just the trip-count check.
+// CHECK: ^[[HEADER]]
+// CHECK: llvm.cond_br %{{.+}}, ^[[BODY:.+]], ^[[EXIT:.+]]
+// Loop body: reuses the entry-block buffer; no fresh alloca here.
+// CHECK: ^[[BODY]]
+// CHECK-NOT: llvm.alloca
+// CHECK: llvm.call @default_cconv_with_extra_fields_in_loop
+// CHECK-SAME: %[[PATCHED_DATA]]
+// CHECK: llvm.br ^[[HEADER]]
+// Loop exit.
+// CHECK: ^[[EXIT]]
+// CHECK: llvm.return
+
+// -----
+
#pipeline_layout = #hal.pipeline.layout<bindings = [
#hal.pipeline.binding<storage_buffer>,
#hal.pipeline.binding<storage_buffer>
