Function to estimate stack size

compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/XCLBinGen.cpp

@@ -6,6 +6,8 @@
 
 #include "XCLBinGen.h"
 
+#include <algorithm>
+#include <charconv>
 #include <filesystem>
 #include <fstream>
 #include <functional>
@@ -60,6 +62,118 @@ using Path = std::filesystem::path;
 namespace mlir::iree_compiler::AMDAIE {
 namespace detail {
 
+// Extract an integer from a string, if possible.
+std::optional<int> safeStoi(std::string_view intString) {
+  size_t start = intString.find_first_not_of(" \t\n\r\f\v");
+  if (start == std::string::npos) return std::nullopt;
+  int value = 0;
+  const char *d0 = intString.data() + start;
+  const char *d1 = intString.data() + intString.size();
+  auto [ptr, ec] = std::from_chars(d0, d1, value);
+  if (ec == std::errc()) return value;
+  return std::nullopt;
+}
+
+// The assembly code for the AIE is of the form (intermediate lines omited):
+// '''
+// .type function_foo,@function
+//  paddb [sp], #224;
+//  function body, includes loads/stores to stack like: 'st  r17, [sp, #-192]'
+//  paddb [sp], #-224;
+//
+// .type function_bar,@function
+//  paddb [sp], #92;
+//  function body, includes loads/stores to stack like: 'st  r17, [sp, #-44]'
+//  paddb [sp], #-92;
+// '''
+//
+// Where above function_foo will use a stack of size 224, and function_bar will
+// use a stack of size 92.
+//
+// We assume that function call structure is as follows:
+// functions with names core_0_0, core_0_1, core_0_2, etc. call into
+// functions with names like generic_matmul_0_outlined. These latter
+// functions do no make any function calls.
+//
+// With the above assumption, the maximum stack size is
+//
+// stack_size = max_{core functions} (stack allocation) +
+//              max_{non-core functions} (stack allocation)
+//
+// Note: this is a temporary lighweight solution while we wait for a
+// better solution in peano: https://github.com/Xilinx/llvm-aie/issues/350
+FailureOr<int> getStackSize(const std::string &aieAssembly) {
+  std::istringstream sstream(aieAssembly);
+  std::string line;
+
+  // Lines like '.type function_name,@function'
+  SmallVector<std::string> funcs;
+
+  // Derived from lines containing substrings like 'paddb [sp], #224;'
+  // There is one vector of ints per function.
+  SmallVector<SmallVector<int>> bumps;
+
+  while (std::getline(sstream, line)) {
+    if (line.find("@function") != std::string::npos) {
+      funcs.push_back(line);
+      bumps.push_back({});
+      continue;
+    }
+    size_t spIndex = line.find("[sp]");
+    if (spIndex == std::string::npos) continue;
+    size_t startIndex = line.rfind(';', spIndex);
+    if (startIndex == std::string::npos) startIndex = 0;
+
+    // I've seen padda, paddb, and one other on strix.
+    size_t paddIndex = line.find("padd", startIndex);
+    if (paddIndex == std::string::npos || paddIndex > spIndex) {
+      llvm::errs() << "Expected 'padd(*) [sp]' in\n" << line << "\n";
+      return failure();
+    }
+    size_t hashIndex = line.find('#', spIndex);
+    if (hashIndex == std::string::npos) {
+      llvm::errs() << "no hash (#) after [sp] in\n" << line << "\n";
+      return failure();
+    }
+
+    // For example, extract 224 from 'paddb [sp], #224;'
+    std::string hashSubstr = line.substr(hashIndex + 1);
+    std::optional<int> maybeNumber = safeStoi(hashSubstr);
+    if (!maybeNumber.has_value()) {
+      llvm::errs() << "failed to get number from \n" << hashSubstr << "\n";
+      return failure();
+    }
+    bumps.back().push_back(maybeNumber.value());
+  }
+
+  // Determine how much the stack pointer is incremented based on the 'padd'
+  // operations extracted above.
+  int maxNonCore{0};
+  int maxCore{0};
+  for (uint32_t i = 0; i < funcs.size(); ++i) {
+    // A function that doesn't use the stack pointer can be ignored.
+    if (bumps[i].empty()) continue;
+
+    // We ignore functions which are not either
+    // 1) sp += x; ...; return
+    // 2) sp += x; ...; sp -= x; return
+    if (bumps[i].size() > 2) {
+      llvm::errs() << "expected no more than 2 sp increments/decrements\n";
+      return failure();
+    }
+    if (bumps[i].size() == 2 && bumps[i][0] + bumps[i][1] != 0) {
+      llvm::errs() << "expected sp increments to cancel out\n";
+      return failure();
+    }
+    if (funcs[i].find("core_") != std::string::npos) {
+      maxCore = std::max(maxCore, bumps[i][0]);
+    } else {
+      maxNonCore = std::max(maxNonCore, bumps[i][0]);
+    }
+  }
+  return maxCore + maxNonCore;
+}
+
 FailureOr<std::vector<std::string>> flagStringToVector(
     const std::string &flags) {
   if (flags.empty()) return std::vector<std::string>{};
@@ -135,7 +249,7 @@ FailureOr<std::vector<std::string>> makePeanoOptArgs(
   // Append the additional flags, unless they conflict with an existing flag,
   // in which case replace the existing flag.
   args.reserve(args.size() + additionalFlags.size());
-  for (const auto &flag : additionalFlags) {
+  for (const std::string &flag : additionalFlags) {
     auto iter = std::find_if(args.begin(), args.end(),
                              std::bind(isContention, _1, flag));
     if (iter == args.end()) {
@@ -235,7 +349,7 @@ FailureOr<Path> findVitis(std::optional<Path> &vitisDir,
   if (!vitisDir) {
     const char *envVitis = ::getenv("VITIS");
     if (!envVitis) {
-      if (auto vpp = sys::findProgramByName("v++")) {
+      if (ErrorOr<std::string> vpp = sys::findProgramByName("v++")) {
         SmallString<64> realVpp;
         std::error_code err = sys::fs::real_path(vpp.get(), realVpp);
         if (!err) {
@@ -429,7 +543,7 @@ LogicalResult runTool(
   {
     std::string prefix{"tmpRunTool"};
     std::string suffix{"Logging"};
-    auto errorCode =
+    std::error_code errorCode =
         llvm::sys::fs::createTemporaryFile(prefix, suffix, temporaryPath);
     if (errorCode) {
       llvm::errs() << "Failed to create temporary file: " << errorCode.message()
@@ -584,7 +698,7 @@ static auto assembleStringUsingPeano =
     std::bind(assembleStringUsing, assembleFileUsingPeano, _1, _2, _3, _4, _5,
               _6, _7, _8, _9);
 
-// Generate the elf files for the core
+// Generate the elf files for the core.
 LogicalResult generateCoreElfFiles(AIE::DeviceOp deviceOp,
                                    const std::string &objFile, Path &tempDir,
                                    bool useChess, bool useChessForUKernel,
@@ -1078,8 +1192,17 @@ void addLowerToLLVMPasses(OpPassManager &pm) {
   pm.addPass(createCSEPass());
 }
 
+FailureOr<int> getStackSizeFromFile(const std::string &asmFileName) {
+  // Read the contents of the file `asmFileName` into a string
+  std::ifstream asmFile(asmFileName);
+  std::string aieAssembly((std::istreambuf_iterator<char>(asmFile)),
+                          std::istreambuf_iterator<char>());
+  return mlir::iree_compiler::AMDAIE::detail::getStackSize(aieAssembly);
+}
+
 LogicalResult generateUnifiedObject(
-    MLIRContext *context, AIE::DeviceOp deviceOp, const std::string &outputFile,
+    MLIRContext *context, AIE::DeviceOp deviceOp,
+    const std::string &outputObjFile, const std::string &outputAsmFile,
     bool printIRBeforeAll, bool printIRAfterAll, bool printIRModuleScope,
     bool timing, bool useChess, bool verbose, Path &tempDir,
     std::optional<Path> vitisDir, const std::string &targetArch, Path &peanoDir,
@@ -1130,7 +1253,7 @@ LogicalResult generateUnifiedObject(
     FailureOr<Path> chessIntrinsicsObjFile = assembleStringUsingChess(
         /*inputFileStr=*/inputLLStr,
         /*inputFileName=*/"input.ll",
-        /*outputFileName=*/outputFile,
+        /*outputFileName=*/outputObjFile,
         /*outputDir=*/tempDir,
         /*extraArgs*/ std::vector<std::string>{},
         /*workDir=*/tempDir,
@@ -1183,15 +1306,51 @@ LogicalResult generateUnifiedObject(
       return failure();
     }
 
-    if (failed(runTool(
-            peanoLLCBin.string(),
-            {OptLLVMIRFile, "-O2", "--march=" + StringRef(targetArch).lower(),
-             "--function-sections", "--filetype=obj", "-o",
-             std::string(outputFile)},
-            verbose))) {
-      llvm::errs() << "Failed to assemble ll with peano\n";
+    std::vector<std::string> llcBaseArguments{
+        OptLLVMIRFile, "-O2", "--march=" + StringRef(targetArch).lower(),
+        "--function-sections"};
+
+    std::vector<std::string> llcObjArguments = llcBaseArguments;
+    llcObjArguments.insert(llcObjArguments.end(),
+                           {"--filetype=obj", "-o", outputObjFile});
+
+    std::vector<std::string> llcAsmArguments = llcBaseArguments;
+    llcAsmArguments.insert(llcAsmArguments.end(),
+                           {"--filetype=asm", "-o", outputAsmFile});
+
+    if (failed(runTool(peanoLLCBin.string(), llcObjArguments, verbose))) {
+      llvm::errs() << "Failed to assemble ll with peano (no .o created)\n";
       return failure();
     }
+
+    if (failed(runTool(peanoLLCBin.string(), llcAsmArguments, verbose))) {
+      llvm::errs() << "Failed to assemble ll with peano (no .s created)\n";
+      return failure();
+    }
+
+    // Stack size check.
+    {
+      FailureOr<int> fromAsm = getStackSizeFromFile(outputAsmFile);
+      if (failed(fromAsm)) {
+        llvm::errs() << "Failed to get stack size from assembly file\n";
+        return failure();
+      }
+      llvm::DenseSet<int> stackSizes;
+      deviceOp->walk([&](AIE::CoreOp coreOp) {
+        stackSizes.insert(coreOp.getStackSize());
+      });
+      if (stackSizes.size() != 1) {
+        llvm::errs() << "Expected a unique assigned stack size\n";
+        return failure();
+      }
+      if (fromAsm.value() > *stackSizes.begin()) {
+        llvm::errs() << "The stack size inferred from the assembly file is "
+                     << fromAsm.value() << " bytes. The assigned stack size is "
+                     << *stackSizes.begin()
+                     << " bytes, which is insufficient. ";
+        return failure();
+      }
+    }
   }
 
   moduleOpCopy->erase();
@@ -1294,11 +1453,12 @@ LogicalResult aie2xclbin(
   if (vitisDirPath) vitisDirPath->make_preferred();
 
   Path unifiedObj = tempDirPath / "input.o";
+  Path unifiedAsm = tempDirPath / "input.s";
   if (failed(generateUnifiedObject(
-          ctx, deviceOp, unifiedObj.string(), printIRBeforeAll, printIRAfterAll,
-          printIRModuleScope, timing, useChess, verbose, tempDirPath,
-          vitisDirPath, targetArch, peanoDirPath, npuVersion,
-          additionalPeanoOptFlags))) {
+          ctx, deviceOp, unifiedObj.string(), unifiedAsm.string(),
+          printIRBeforeAll, printIRAfterAll, printIRModuleScope, timing,
+          useChess, verbose, tempDirPath, vitisDirPath, targetArch,
+          peanoDirPath, npuVersion, additionalPeanoOptFlags))) {
     llvm::errs() << "Failed to generate unified object\n";
     return failure();
   }

compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/XCLBinGen.h

@@ -32,9 +32,20 @@ mlir::LogicalResult emitNpuInstructions(xilinx::AIE::DeviceOp deviceOp,
                                         const std::string &outputNPU);
 
 namespace detail {
+
 FailureOr<std::vector<std::string>> flagStringToVector(
     const std::string &flags);
+
 FailureOr<std::vector<std::string>> makePeanoOptArgs(
     const std::vector<std::string> &additionalPeanoOptFlags);
+
+/// An exception-free version of std::stoi, using C++17's std::from_chars.
+std::optional<int> safeStoi(std::string_view intString);
+
+/// Get the stack size by analysing AIE assembly code. This function is not
+/// guaranteed to work for all AIE assembly code, and should be used with
+/// caution.
+FailureOr<int> getStackSize(const std::string &aieAssembly);
+
 }  // namespace detail
 }  // namespace mlir::iree_compiler::AMDAIE

compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/test/AMDAIEXCLBinGenTest.cpp

@@ -44,6 +44,48 @@ TEST(XCLBinGenTest, makePeanoOptArgs) {
   EXPECT_FALSE(succeeded(maybeAdditionalFlagsWithoutSemis));
 }
 
+TEST(XCLBinGenTest, SafeStoi0) {
+  // Basically this function
+  // (1) strips leading whitespace, then
+  // (2) finds the largest valid integer in the string.
+  EXPECT_TRUE(detail::safeStoi("   123; 4.123") == 123);
+  EXPECT_TRUE(detail::safeStoi("   -3; ") == -3);
+}
+
+TEST(XCLBinGenTest, GetStackSize0) {
+  // Intermediate lines of assembly omitted for brevity.
+  {
+    std::string asmStr = R"(
+ .type generic_matmul_0_outlined,@function
+ .type core_0_2,@function
+ nopa ;  paddb [sp], #224;  nopxm ;  nops
+ paddb [sp], #-224                     //  Delay Slot 1
+)";
+    EXPECT_TRUE(detail::getStackSize(asmStr) == 224);
+  }
+
+  {
+    std::string asmStr = R"(
+ .type core_0_2,@function
+ nopa ;  paddb [sp], #96;  nopxm ;  nops
+)";
+    EXPECT_TRUE(detail::getStackSize(asmStr) == 96);
+  }
+
+  {
+    std::string asmStr = R"(
+ .type generic_matmul_0_outlined,@function
+ nopa ;  padda [sp], #20;  nopxm ;  nops
+ padda [sp], #-20                     //  Delay Slot 1
+ .type core_0_2,@function
+ nopa ;  padda [sp], #224;  nopxm ;  nops
+ .type core_0_3,@function
+ nopa ;  padda [sp], #224;  nopxm ;  nops
+)";
+    EXPECT_TRUE(detail::getStackSize(asmStr) == 224 + 20);
+  }
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {