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

[Min-Yih Hsu via llvm-commits]

https://github.com/mshockwave updated https://github.com/llvm/llvm-project/pull/79636

>From 1567f28d1646ad8127bab1aeecd686a6cd018b0b Mon Sep 17 00:00:00 2001
From: Min Hsu <min.hsu at sifive.com>
Date: Fri, 26 Jan 2024 11:00:33 -0800
Subject: [PATCH 1/2] [Exegesis] Do not assume the size and layout of the
 assembled snippet

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.
---
 llvm/tools/llvm-exegesis/lib/Assembler.cpp | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/llvm/tools/llvm-exegesis/lib/Assembler.cpp b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
index 9f03a4e3a5a6ff4..366971ed6a546dc 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.cpp
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
@@ -365,11 +365,18 @@ 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;
+    // Suppress the error.
+    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)

>From 43b01aedc368ccc4da53e207996ae4623ff71882 Mon Sep 17 00:00:00 2001
From: Min Hsu <min.hsu at sifive.com>
Date: Fri, 26 Jan 2024 15:07:46 -0800
Subject: [PATCH 2/2] fixup! [Exegesis] Do not assume the size and layout of
 the assembled snippet

---
 llvm/tools/llvm-exegesis/lib/Assembler.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/tools/llvm-exegesis/lib/Assembler.cpp b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
index 366971ed6a546dc..307b951f4a8496e 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.cpp
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
@@ -370,7 +370,9 @@ Expected<ExecutableFunction> ExecutableFunction::create(
     auto SymbolName = Pair.first.getName();
     if (SymbolName)
       return *SymbolName == FunctionID;
-    // Suppress the error.
+    // 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;
   });