1 files changed, 15 insertions, 5 deletions
diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp
index 1dd413a2912..c7a0c8470bb 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp
@@ -32,7 +32,8 @@ struct YamlContext {
   YamlContext(const exegesis::LLVMState &State)
       : State(&State), ErrorStream(LastError),
         OpcodeNameToOpcodeIdx(
-            generateOpcodeNameToOpcodeIdxMapping(State.getInstrInfo())) {}
+            generateOpcodeNameToOpcodeIdxMapping(State.getInstrInfo())),
+        RegNameToRegNo(generateRegNameToRegNoMapping(State.getRegInfo())) {}
 
   static llvm::StringMap<unsigned>
   generateOpcodeNameToOpcodeIdxMapping(const llvm::MCInstrInfo &InstrInfo) {
@@ -43,6 +44,15 @@ struct YamlContext {
     return Map;
   };
 
+  llvm::StringMap<unsigned>
+  generateRegNameToRegNoMapping(const llvm::MCRegisterInfo &RegInfo) {
+    llvm::StringMap<unsigned> Map(RegInfo.getNumRegs());
+    for (unsigned I = 0, E = RegInfo.getNumRegs(); I < E; ++I)
+      Map[RegInfo.getName(I)] = I;
+    assert(Map.size() == RegInfo.getNumRegs() && "Size prediction failed");
+    return Map;
+  };
+
   void serializeMCInst(const llvm::MCInst &MCInst, llvm::raw_ostream &OS) {
     OS << getInstrName(MCInst.getOpcode());
     for (const auto &Op : MCInst) {
@@ -80,10 +90,9 @@ struct YamlContext {
   }
 
   unsigned getRegNo(llvm::StringRef RegName) {
-    const llvm::MCRegisterInfo &RegInfo = State->getRegInfo();
-    for (unsigned E = RegInfo.getNumRegs(), I = 0; I < E; ++I)
-      if (RegInfo.getName(I) == RegName)
-        return I;
+    auto Iter = RegNameToRegNo.find(RegName);
+    if (Iter != RegNameToRegNo.end())
+      return Iter->second;
     ErrorStream << "No register with name " << RegName;
     return 0;
   }
@@ -159,6 +168,7 @@ private:
   std::string LastError;
   llvm::raw_string_ostream ErrorStream;
   const llvm::StringMap<unsigned> OpcodeNameToOpcodeIdx;
+  const llvm::StringMap<unsigned> RegNameToRegNo;
 };
 } // namespace