[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 created https://github.com/llvm/llvm-project/pull/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.

>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] [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)