Enable checking stack allocation before converting to LLVM dialect. (#…

…8976)

The flag is not removed because sometimes we might want to bypass the check.

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

iree/compiler/Codegen/LLVMCPU/Passes.cpp

@@ -33,7 +33,7 @@ static llvm::cl::opt<bool> 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

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
+}
+}