diff --git a/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp b/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp index d45b87a39c36..7173ab936d9f 100644 --- a/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp +++ b/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp @@ -21,13 +21,25 @@ struct LLVMCPUCheckIRBeforeLLVMConversionPass 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(); + 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> diff --git a/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/iree/compiler/Codegen/LLVMCPU/Passes.cpp index ee7cc49567c4..830366e29f6a 100644 --- a/iree/compiler/Codegen/LLVMCPU/Passes.cpp +++ b/iree/compiler/Codegen/LLVMCPU/Passes.cpp @@ -33,7 +33,7 @@ static llvm::cl::opt clCheckIRBeforeLLVMConversion( "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 b646973db9d3..2765093decb3 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 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 +} +}