[LLVMCPU] Set stack allocation limit as a HAL target property (#20209)

As long as the limit is defined arbitrarily as a pass-specific CLI
option, tiling/bufferization passes have no way of accounting for it.
This change would be the first step towards addressing compilation
failures in `LLVMCPUCheckIRBeforeLLVMConversionPass` systemically before
further tuning of the tiling/fusion logic gets done for particular
benchmarks.

A basic LIT test is also added to make sure that the limit can still be
overriden through a CLI option. This can be expanded further to test
other default attribute values.

Signed-off-by: Artem Gindinson <[email protected]>

---------

Signed-off-by: Artem Gindinson <[email protected]>

Author: AGindinson

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<llvm::cl::PowerOf2ByteSize>(
+      "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.value;
   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 int64_t DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES = 32768;
   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;
+  int64_t 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;
+  llvm::cl::PowerOf2ByteSize 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/plugins/target/LLVMCPU/test/BUILD.bazel

@@ -16,6 +16,7 @@ iree_lit_test_suite(
     name = "lit",
     srcs = enforce_glob(
         [
+            "hal_target_device_attributes.mlir",
             "materialize_homogeneous_encodings.mlir",
             "smoketest_embedded.mlir",
             "smoketest_system.mlir",
@@ -24,8 +25,10 @@ iree_lit_test_suite(
     ),
     cfg = "//compiler:lit.cfg.py",
     tools = [
+        "//tools:iree-compile",
         "//tools:iree-opt",
         "@llvm-project//lld",
         "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
     ],
 )

compiler/plugins/target/LLVMCPU/test/CMakeLists.txt

@@ -14,13 +14,16 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
+    "hal_target_device_attributes.mlir"
     "materialize_homogeneous_encodings.mlir"
     "smoketest_embedded.mlir"
     "smoketest_system.mlir"
   TOOLS
     ${IREE_LLD_TARGET}
     FileCheck
+    iree-compile
     iree-opt
+    not
 )
 
 ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

compiler/plugins/target/LLVMCPU/test/hal_target_device_attributes.mlir

@@ -0,0 +1,39 @@
+// 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 = 32768 : 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
+
+// RUN: not 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=64266 \
+// RUN: 2>&1 | FileCheck %s --check-prefix=CHECK-INCORRECT-OPT-STACK-VALUE
+//
+// CHECK-INCORRECT-OPT-STACK-VALUE: for the --iree-llvmcpu-stack-allocation-limit option: '64266' value not a power-of-two
+
+module {
+  util.func public @foo(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+    util.return %arg0 : tensor<?xf32>
+  }
+}

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(
@@ -38,13 +34,26 @@ struct LLVMCPUCheckIRBeforeLLVMConversionPass
 };
 } // namespace
 
-/// Returns success if the cummulative stack allocation size is less than the
-/// limit set by clMaxAllocationSizeInBytes.
+/// Returns success if the cumulative stack allocation size is less than the
+/// limit set through --iree-llvmcpu-stack-allocation-limit (or the default
+/// defined for HAL LLVMCPU target).
 static LogicalResult
 checkStackAllocationSize(mlir::FunctionOpInterface funcOp) {
   if (funcOp.getFunctionBody().empty())
     return success();
 
+  // In rare cases where the attribute is not present in the module, a value of
+  // 32KB will be taken.
+  unsigned maxAllocationSizeInBytes = 32 * 1024;
+  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 +100,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();
 }