refactor(iree): set CPU stack allocation limit as an LLVM target option

compiler/plugins/target/LLVMCPU/LLVMTargetOptions.cpp

@@ -139,6 +139,7 @@ void LLVMTarget::storeToConfigAttrs(MLIRContext *context,
   if (vectorWidthInBytes != DEFAULT_VECTOR_WIDTH_IN_BYTES) {
     addInt64("native_vector_size", vectorWidthInBytes);
   }
+  addInt64("max_stack_allocation_size", maxStackAllocSizeInBytes);
   if (linkEmbedded != DEFAULT_LINK_EMBEDDED) {
     addBool("link_embedded", linkEmbedded);
   }
@@ -581,6 +582,11 @@ void LLVMCPUTargetCLOptions::bindOptions(OptionsBinder &binder) {
                        targetVectorWidthInBytes, llvm::cl::cat(category),
                        llvm::cl::desc("Overrides the native vector register "
                                       "width (in bytes) of the target."));
+  binder.opt<unsigned>(
+      "iree-llvmcpu-stack-allocation-limit", targetMaxStackAllocSizeInBytes,
+      llvm::cl::cat(category),
+      llvm::cl::desc(
+          "Maximum allowed stack allocation size for LLVM CPU in bytes"));
   binder.opt<std::string>(
       "iree-llvmcpu-enable-ukernels", enableUkernels, llvm::cl::cat(category),
       llvm::cl::desc("Enables ukernels in the llvmcpu backend. May be "
@@ -634,6 +640,7 @@ LLVMTargetOptions LLVMCPUTargetCLOptions::getTargetOptions() {
   target.llvmTargetOptions.FloatABIType = targetFloatABI;
   target.dataLayout = targetDataLayout;
   target.vectorWidthInBytes = targetVectorWidthInBytes;
+  target.maxStackAllocSizeInBytes = targetMaxStackAllocSizeInBytes;
   target.ukernels = enableUkernels;
   target.linkUkernelBitcode = linkUKernelBitcode;
 

compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h

@@ -33,6 +33,7 @@ enum class SanitizerKind {
 struct LLVMTarget {
   static constexpr const char *DEFAULT_DATA_LAYOUT = "";
   static constexpr int64_t DEFAULT_VECTOR_WIDTH_IN_BYTES = 0;
+  static constexpr unsigned DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES = 32678;
   static constexpr bool DEFAULT_LINK_EMBEDDED = true;
   static constexpr bool DEFAULT_DEBUG_SYMBOLS = true;
   static constexpr SanitizerKind DEFAULT_SANITIZER_KIND = SanitizerKind::kNone;
@@ -88,6 +89,7 @@ struct LLVMTarget {
   std::string dataLayout = DEFAULT_DATA_LAYOUT;
   // Overrides the vector width (in bytes) of the target.
   int64_t vectorWidthInBytes = DEFAULT_VECTOR_WIDTH_IN_BYTES;
+  unsigned maxStackAllocSizeInBytes = DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES;
 
   llvm::PipelineTuningOptions pipelineTuningOptions;
   // Optimization level to be used by the LLVM optimizer (middle-end).
@@ -194,6 +196,8 @@ struct LLVMCPUTargetCLOptions {
   llvm::FloatABI::ABIType targetFloatABI = LLVMTarget::DEFAULT_FLOAT_ABI;
   std::string targetDataLayout = LLVMTarget::DEFAULT_DATA_LAYOUT;
   unsigned targetVectorWidthInBytes = LLVMTarget::DEFAULT_VECTOR_WIDTH_IN_BYTES;
+  unsigned targetMaxStackAllocSizeInBytes =
+      LLVMTarget::DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES;
   std::string enableUkernels = LLVMTarget::DEFAULT_ENABLE_UKERNELS;
   bool linkUKernelBitcode = LLVMTarget::DEFAULT_LINK_UKERNEL_BITCODE;
   bool listTargets; // Ignored - used with llvm::cl::ValueDisallowed.

compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp

@@ -5,6 +5,7 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Codegen/LLVMCPU/Passes.h"
+#include "iree/compiler/Codegen/Utils/Utils.h"
 #include "llvm/Support/CommandLine.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
@@ -16,11 +17,6 @@ namespace mlir::iree_compiler {
 #define GEN_PASS_DEF_LLVMCPUCHECKIRBEFORELLVMCONVERSIONPASS
 #include "iree/compiler/Codegen/LLVMCPU/Passes.h.inc"
 
-static llvm::cl::opt<int> clMaxAllocationSizeInBytes(
-    "iree-llvmcpu-stack-allocation-limit",
-    llvm::cl::desc("maximum allowed stack allocation size in bytes"),
-    llvm::cl::init(32768));
-
 static llvm::cl::opt<unsigned> clAssumedVscaleValue(
     "iree-llvmcpu-stack-allocation-assumed-vscale",
     llvm::cl::desc(
@@ -45,6 +41,16 @@ checkStackAllocationSize(mlir::FunctionOpInterface funcOp) {
   if (funcOp.getFunctionBody().empty())
     return success();
 
+  unsigned maxAllocationSizeInBytes = 32768;
+  auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(funcOp);
+  if (targetAttr) {
+    auto nativeAllocationSizeAttr =
+        getConfigIntegerAttr(targetAttr, "max_stack_allocation_size");
+    if (nativeAllocationSizeAttr) {
+      maxAllocationSizeInBytes = nativeAllocationSizeAttr->getInt();
+    }
+  }
+
   SmallVector<memref::AllocaOp> allocaOps;
   funcOp.walk(
       [&](memref::AllocaOp allocaOp) { allocaOps.push_back(allocaOp); });
@@ -91,10 +97,10 @@ checkStackAllocationSize(mlir::FunctionOpInterface funcOp) {
     }
     cumSize += allocaSize / 8;
   }
-  if (cumSize > clMaxAllocationSizeInBytes) {
+  if (cumSize > maxAllocationSizeInBytes) {
     return funcOp.emitOpError("exceeded stack allocation limit of ")
-           << clMaxAllocationSizeInBytes.getValue()
-           << " bytes for function. Got " << cumSize << " bytes";
+           << maxAllocationSizeInBytes << " bytes for function. Got "
+           << cumSize << " bytes";
   }
   return success();
 }

compiler/src/iree/compiler/Codegen/LLVMCPU/test/hal_target_device_attributes.mlir

@@ -0,0 +1,33 @@
+// This test aims to check default HAL properties for LLVM CPU target, and
+// whether CLI options modify the values correctly.
+
+// TODO: Expand the test for more CLI configurations, e.g. different target triples
+
+// RUN: iree-compile --compile-to=preprocessing --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-triple=x86_64-linux-gnu %s \
+// RUN: | FileCheck %s --check-prefix=CHECK-X86-DEFAULT
+//
+// CHECK-X86-DEFAULT: module attributes {stream.affinity.default = #hal.device.affinity<@__device_0>} {
+// CHECK-X86-DEFAULT-NEXT: util.global private @__device_0 = #hal.device.target<"local",
+// CHECK-X86-DEFAULT-SAME: [#hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "", cpu_features = ""
+// CHECK-X86-DEFAULT-SAME: max_stack_allocation_size = 32678 : i64
+// CHECK-X86-DEFAULT-SAME: native_vector_size = 16 : i64
+// CHECK-X86-DEFAULT-SAME: target_triple = "x86_64-unknown-unknown-eabi-elf"
+// CHECK-X86-DEFAULT-SAME: }>]> : !hal.device
+
+// RUN: iree-compile --compile-to=preprocessing --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-triple=x86_64-linux-gnu %s \
+// RUN:              --iree-llvmcpu-stack-allocation-limit=65536 \
+// RUN: | FileCheck %s --check-prefix=CHECK-STACK-VALUE
+//
+// CHECK-STACK-VALUE: module attributes {stream.affinity.default = #hal.device.affinity<@__device_0>} {
+// CHECK-STACK-VALUE-NEXT: util.global private @__device_0 = #hal.device.target<"local",
+// CHECK-STACK-VALUE-SAME: [#hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "", cpu_features = ""
+//
+// CHECK-STACK-VALUE-SAME: max_stack_allocation_size = 65536 : i64
+//
+// CHECK-STACK-VALUE-SAME: }>]> : !hal.device
+
+module {
+  util.func public @foo(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+    util.return %arg0 : tensor<?xf32>
+  }
+}