Enable checking stack allocation before converting to LLVM dialect. (#8976)
The flag is not removed because sometimes we might want to bypass the check.
diff --git a/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp b/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp
index d45b87a..7173ab9 100644
--- a/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp
+++ b/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp
@@ -21,13 +21,25 @@
void LLVMCPUCheckIRBeforeLLVMConversionPass::runOnOperation() {
auto moduleOp = getOperation();
- // For now only check that there are no stack allocations.
- auto walkResult = moduleOp.walk([](memref::AllocaOp allocaOp) -> WalkResult {
- return allocaOp.emitOpError("expected no static allocations");
+ int64_t bits = 0;
+ auto walkResult = moduleOp.walk([&](memref::AllocaOp allocaOp) -> WalkResult {
+ auto type = allocaOp.getType().cast<ShapedType>();
+ if (!type.hasStaticShape()) {
+ return allocaOp.emitOpError(
+ "expected no stack allocations with dynamic shapes");
+ }
+ bits += type.getSizeInBits();
+ return WalkResult::advance();
});
if (walkResult.wasInterrupted()) {
return signalPassFailure();
}
+ constexpr int k16KBInBits = 16 * 1024 * 8;
+ if (bits >= k16KBInBits) {
+ moduleOp.emitOpError(
+ "expected total size of stack allocation is smaller than 16 KB");
+ return signalPassFailure();
+ }
}
std::unique_ptr<OperationPass<ModuleOp>>
diff --git a/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index ee7cc49..830366e 100644
--- a/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -33,7 +33,7 @@
"iree-codegen-check-ir-before-llvm-conversion",
llvm::cl::desc("Runs the pass to check the IR generated from LLVMCPU "
"before conversion to LLVM IR"),
- llvm::cl::init(false));
+ llvm::cl::init(true));
//===---------------------------------------------------------------------===//
// Default allocation functions for CPU backend
diff --git a/iree/compiler/Codegen/LLVMCPU/test/check_ir_before_llvm_conversion.mlir b/iree/compiler/Codegen/LLVMCPU/test/check_ir_before_llvm_conversion.mlir
index b646973..2765093 100644
--- a/iree/compiler/Codegen/LLVMCPU/test/check_ir_before_llvm_conversion.mlir
+++ b/iree/compiler/Codegen/LLVMCPU/test/check_ir_before_llvm_conversion.mlir
@@ -1,7 +1,19 @@
// RUN: iree-opt -iree-llvmcpu-check-ir-before-llvm-conversion %s -verify-diagnostics -split-input-file
-func.func @no_static_allocas(%arg0: index) {
- // expected-error @+1 {{expected no static allocations}}
+module {
+func.func @no_dynamic_allocas(%arg0: index) {
+ // expected-error @+1 {{expected no stack allocations with dynamic shapes}}
%0 = memref.alloca(%arg0) : memref<?xf32>
return
}
+}
+
+// -----
+
+// expected-error @+1 {{expected total size of stack allocation is smaller than 16 KB}}
+module {
+func.func @big_allocas(%arg0: index) {
+ %0 = memref.alloca() : memref<65536xi32>
+ return
+}
+}
