[llvm] 3bece3d - [Exegesis] Do not assume the size and layout of the assembled snippet (#79636)

[via llvm-commits]

Author: Min-Yih Hsu
Date: 2024-01-26T15:47:26-08:00
New Revision: 3bece3d72d1323af6b2418b13281e66b583f25b3

URL: https://github.com/llvm/llvm-project/commit/3bece3d72d1323af6b2418b13281e66b583f25b3
DIFF: https://github.com/llvm/llvm-project/commit/3bece3d72d1323af6b2418b13281e66b583f25b3.diff

LOG: [Exegesis] Do not assume the size and layout of the assembled snippet (#79636)

Currently llvm-exegesis assumes that there will only be 3 symbols in the
snippet object, in which the benchmarking function 'foo' is always the
last symbol.

These assumptions do not hold for object file formats of other targets
we support downstream. I think it would be more ideal to generalize this
part of the logics into a simple search on all symbols, as proposed by
this patch.

Added: 
    

Modified: 
    llvm/tools/llvm-exegesis/lib/Assembler.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/tools/llvm-exegesis/lib/Assembler.cpp b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
index 9f03a4e3a5a6ff4..307b951f4a8496e 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.cpp
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
@@ -365,11 +365,20 @@ Expected<ExecutableFunction> ExecutableFunction::create(
 
   auto SymbolSizes = object::computeSymbolSizes(*ObjectFileHolder.getBinary());
   // Get the size of the function that we want to call into (with the name of
-  // FunctionID). This should always be the third symbol returned by
-  // calculateSymbolSizes.
-  assert(SymbolSizes.size() == 3);
-  assert(cantFail(std::get<0>(SymbolSizes[2]).getName()) == FunctionID);
-  uintptr_t CodeSize = std::get<1>(SymbolSizes[2]);
+  // FunctionID).
+  auto SymbolIt = llvm::find_if(SymbolSizes, [&](const auto &Pair) {
+    auto SymbolName = Pair.first.getName();
+    if (SymbolName)
+      return *SymbolName == FunctionID;
+    // We should always succeed in finding the FunctionID, hence we suppress
+    // the error here and assert later on the search result, rather than
+    // propagating the Expected<> error back to the caller.
+    llvm::consumeError(SymbolName.takeError());
+    return false;
+  });
+  assert(SymbolIt != SymbolSizes.end() &&
+         "Cannot find the symbol for FunctionID");
+  uintptr_t CodeSize = SymbolIt->second;
 
   auto EJITOrErr = orc::LLJITBuilder().create();
   if (!EJITOrErr)