fix(iree): adjust stack allocation patches per upstream review (#37)

- Address formatting issues
- Move the LIT test, register with Bazel & CMake
- Use `PowerOf2ByteSize` for the CLI option, add negative LIT check
- Improve readability in the verification pass

Author: AGindinson

Diff for: compiler/plugins/target/LLVMCPU/LLVMTargetOptions.cpp

@@ -582,7 +582,7 @@ 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>(
+  binder.opt<llvm::cl::PowerOf2ByteSize>(
       "iree-llvmcpu-stack-allocation-limit", targetMaxStackAllocSizeInBytes,
       llvm::cl::cat(category),
       llvm::cl::desc(
@@ -640,7 +640,7 @@ LLVMTargetOptions LLVMCPUTargetCLOptions::getTargetOptions() {
   target.llvmTargetOptions.FloatABIType = targetFloatABI;
   target.dataLayout = targetDataLayout;
   target.vectorWidthInBytes = targetVectorWidthInBytes;
-  target.maxStackAllocSizeInBytes = targetMaxStackAllocSizeInBytes;
+  target.maxStackAllocSizeInBytes = targetMaxStackAllocSizeInBytes.value;
   target.ukernels = enableUkernels;
   target.linkUkernelBitcode = linkUKernelBitcode;
 

Diff for: compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h

@@ -33,7 +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 = 32768;
+  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;
@@ -89,7 +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;
+  int64_t maxStackAllocSizeInBytes = DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES;
 
   llvm::PipelineTuningOptions pipelineTuningOptions;
   // Optimization level to be used by the LLVM optimizer (middle-end).
@@ -196,7 +196,7 @@ 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 =
+  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;

Diff for: 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",
     ],
 )

Diff for: 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 ###

Diff for: compiler/src/iree/compiler/Codegen/LLVMCPU/test/hal_target_device_attributes.mlir renamed to compiler/plugins/target/LLVMCPU/test/hal_target_device_attributes.mlir

@@ -26,6 +26,12 @@
 //
 // 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>

Diff for: compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp

@@ -34,14 +34,17 @@ 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();
 
-  unsigned maxAllocationSizeInBytes = 32768;
+  // 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 =
@@ -99,8 +102,8 @@ checkStackAllocationSize(mlir::FunctionOpInterface funcOp) {
   }
   if (cumSize > maxAllocationSizeInBytes) {
     return funcOp.emitOpError("exceeded stack allocation limit of ")
-           << maxAllocationSizeInBytes << " bytes for function. Got "
-           << cumSize << " bytes";
+           << maxAllocationSizeInBytes << " bytes for function. Got " << cumSize
+           << " bytes";
   }
   return success();
 }