[llvm-exegesis] Refactor getting register number from name to LLVMState (#107895)

This patch refactors the procedure of getting the register number from a
register name to LLVMState rather than having individual users get the
values themselves by getting a reference to the map from LLVMState. This
is primarily intended to make some downstream usage in Gematria simpler,
but also cleans up a little bit upstream by pulling the actual map
searching out and just leaving error handling to the clients.

The original getter is left to enable downstream migration in Gematria,
particularly before it gets imported into google internal.

Author: boomanaiden154

llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp

@@ -34,8 +34,7 @@ namespace {
 struct YamlContext {
   YamlContext(const exegesis::LLVMState &State)
       : State(&State), ErrorStream(LastError),
-        OpcodeNameToOpcodeIdx(State.getOpcodeNameToOpcodeIdxMapping()),
-        RegNameToRegNo(State.getRegNameToRegNoMapping()) {}
+        OpcodeNameToOpcodeIdx(State.getOpcodeNameToOpcodeIdxMapping()) {}
 
   void serializeMCInst(const MCInst &MCInst, raw_ostream &OS) {
     OS << getInstrName(MCInst.getOpcode());
@@ -77,11 +76,11 @@ struct YamlContext {
   }
 
   std::optional<unsigned> getRegNo(StringRef RegName) {
-    auto Iter = RegNameToRegNo.find(RegName);
-    if (Iter != RegNameToRegNo.end())
-      return Iter->second;
-    ErrorStream << "No register with name '" << RegName << "'\n";
-    return std::nullopt;
+    std::optional<MCRegister> RegisterNumber =
+        State->getRegisterNumberFromName(RegName);
+    if (!RegisterNumber.has_value())
+      ErrorStream << "No register with name '" << RegName << "'\n";
+    return RegisterNumber;
   }
 
 private:
@@ -154,7 +153,6 @@ struct YamlContext {
   std::string LastError;
   raw_string_ostream ErrorStream;
   const DenseMap<StringRef, unsigned> &OpcodeNameToOpcodeIdx;
-  const DenseMap<StringRef, MCRegister> &RegNameToRegNo;
 };
 } // namespace
 
@@ -231,10 +229,8 @@ template <> struct MappingTraits<exegesis::BenchmarkMeasure> {
   static const bool flow = true;
 };
 
-template <>
-struct ScalarEnumerationTraits<exegesis::Benchmark::ModeE> {
-  static void enumeration(IO &Io,
-                          exegesis::Benchmark::ModeE &Value) {
+template <> struct ScalarEnumerationTraits<exegesis::Benchmark::ModeE> {
+  static void enumeration(IO &Io, exegesis::Benchmark::ModeE &Value) {
     Io.enumCase(Value, "", exegesis::Benchmark::Unknown);
     Io.enumCase(Value, "latency", exegesis::Benchmark::Latency);
     Io.enumCase(Value, "uops", exegesis::Benchmark::Uops);
@@ -282,8 +278,7 @@ template <> struct ScalarTraits<exegesis::RegisterValue> {
   static const bool flow = true;
 };
 
-template <>
-struct MappingContextTraits<exegesis::BenchmarkKey, YamlContext> {
+template <> struct MappingContextTraits<exegesis::BenchmarkKey, YamlContext> {
   static void mapping(IO &Io, exegesis::BenchmarkKey &Obj,
                       YamlContext &Context) {
     Io.setContext(&Context);
@@ -293,8 +288,7 @@ struct MappingContextTraits<exegesis::BenchmarkKey, YamlContext> {
   }
 };
 
-template <>
-struct MappingContextTraits<exegesis::Benchmark, YamlContext> {
+template <> struct MappingContextTraits<exegesis::Benchmark, YamlContext> {
   struct NormalizedBinary {
     NormalizedBinary(IO &io) {}
     NormalizedBinary(IO &, std::vector<uint8_t> &Data) : Binary(Data) {}
@@ -311,8 +305,7 @@ struct MappingContextTraits<exegesis::Benchmark, YamlContext> {
     BinaryRef Binary;
   };
 
-  static void mapping(IO &Io, exegesis::Benchmark &Obj,
-                      YamlContext &Context) {
+  static void mapping(IO &Io, exegesis::Benchmark &Obj, YamlContext &Context) {
     Io.mapRequired("mode", Obj.Mode);
     Io.mapRequired("key", Obj.Key, Context);
     Io.mapRequired("cpu_name", Obj.CpuName);
@@ -339,8 +332,7 @@ struct MappingContextTraits<exegesis::Benchmark, YamlContext> {
 };
 
 template <> struct MappingTraits<exegesis::Benchmark::TripleAndCpu> {
-  static void mapping(IO &Io,
-                      exegesis::Benchmark::TripleAndCpu &Obj) {
+  static void mapping(IO &Io, exegesis::Benchmark::TripleAndCpu &Obj) {
     assert(!Io.outputting() && "can only read TripleAndCpu");
     // Read triple.
     Io.mapRequired("llvm_triple", Obj.LLVMTriple);
@@ -357,8 +349,7 @@ Expected<std::set<Benchmark::TripleAndCpu>>
 Benchmark::readTriplesAndCpusFromYamls(MemoryBufferRef Buffer) {
   // We're only mapping a field, drop other fields and silence the corresponding
   // warnings.
-  yaml::Input Yin(
-      Buffer, nullptr, +[](const SMDiagnostic &, void *Context) {});
+  yaml::Input Yin(Buffer, nullptr, +[](const SMDiagnostic &, void *Context) {});
   Yin.setAllowUnknownKeys(true);
   std::set<TripleAndCpu> Result;
   yaml::EmptyContext Context;
@@ -373,8 +364,8 @@ Benchmark::readTriplesAndCpusFromYamls(MemoryBufferRef Buffer) {
   return Result;
 }
 
-Expected<Benchmark>
-Benchmark::readYaml(const LLVMState &State, MemoryBufferRef Buffer) {
+Expected<Benchmark> Benchmark::readYaml(const LLVMState &State,
+                                        MemoryBufferRef Buffer) {
   yaml::Input Yin(Buffer);
   YamlContext Context(State);
   Benchmark Benchmark;
@@ -385,9 +376,8 @@ Benchmark::readYaml(const LLVMState &State, MemoryBufferRef Buffer) {
   return std::move(Benchmark);
 }
 
-Expected<std::vector<Benchmark>>
-Benchmark::readYamls(const LLVMState &State,
-                                MemoryBufferRef Buffer) {
+Expected<std::vector<Benchmark>> Benchmark::readYamls(const LLVMState &State,
+                                                      MemoryBufferRef Buffer) {
   yaml::Input Yin(Buffer);
   YamlContext Context(State);
   std::vector<Benchmark> Benchmarks;
@@ -403,8 +393,7 @@ Benchmark::readYamls(const LLVMState &State,
   return std::move(Benchmarks);
 }
 
-Error Benchmark::writeYamlTo(const LLVMState &State,
-                                        raw_ostream &OS) {
+Error Benchmark::writeYamlTo(const LLVMState &State, raw_ostream &OS) {
   auto Cleanup = make_scope_exit([&] { OS.flush(); });
   yaml::Output Yout(OS, nullptr /*Ctx*/, 200 /*WrapColumn*/);
   YamlContext Context(State);
@@ -416,8 +405,7 @@ Error Benchmark::writeYamlTo(const LLVMState &State,
   return Error::success();
 }
 
-Error Benchmark::readYamlFrom(const LLVMState &State,
-                                         StringRef InputContent) {
+Error Benchmark::readYamlFrom(const LLVMState &State, StringRef InputContent) {
   yaml::Input Yin(InputContent);
   YamlContext Context(State);
   if (Yin.setCurrentDocument())

llvm/tools/llvm-exegesis/lib/LlvmState.cpp

@@ -100,6 +100,14 @@ std::unique_ptr<LLVMTargetMachine> LLVMState::createTargetMachine() const {
           Reloc::Model::Static)));
 }
 
+std::optional<MCRegister>
+LLVMState::getRegisterNumberFromName(StringRef RegisterName) const {
+  auto RegisterIt = RegNameToRegNoMapping->find(RegisterName);
+  if (RegisterIt == RegNameToRegNoMapping->end())
+    return std::nullopt;
+  return RegisterIt->second;
+}
+
 std::unique_ptr<const DenseMap<StringRef, unsigned>>
 LLVMState::createOpcodeNameToOpcodeIdxMapping() const {
   const MCInstrInfo &InstrInfo = getInstrInfo();

llvm/tools/llvm-exegesis/lib/LlvmState.h

@@ -76,11 +76,17 @@ class LLVMState {
     return *OpcodeNameToOpcodeIdxMapping;
   };
 
+  // TODO(boomanaiden154): We are keeping this getter around to enable internal
+  // migration to getRegisterNumberFromName. Once that is complete and
+  // the changes have been pulled, we can remove this.
   const DenseMap<StringRef, MCRegister> &getRegNameToRegNoMapping() const {
     assert(RegNameToRegNoMapping);
     return *RegNameToRegNoMapping;
   }
 
+  std::optional<MCRegister>
+  getRegisterNumberFromName(StringRef RegisterName) const;
+
 private:
   std::unique_ptr<const DenseMap<StringRef, unsigned>>
   createOpcodeNameToOpcodeIdxMapping() const;

llvm/tools/llvm-exegesis/lib/SnippetFile.cpp

@@ -9,13 +9,15 @@
 #include "SnippetFile.h"
 #include "BenchmarkRunner.h"
 #include "Error.h"
+#include "LlvmState.h"
 #include "Target.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCParser/MCAsmLexer.h"
 #include "llvm/MC/MCParser/MCAsmParser.h"
 #include "llvm/MC/MCParser/MCTargetAsmParser.h"
+#include "llvm/MC/MCRegister.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/TargetRegistry.h"
@@ -35,10 +37,9 @@ namespace {
 // An MCStreamer that reads a BenchmarkCode definition from a file.
 class BenchmarkCodeStreamer : public MCStreamer, public AsmCommentConsumer {
 public:
-  explicit BenchmarkCodeStreamer(
-      MCContext *Context, const DenseMap<StringRef, MCRegister> &RegNameToRegNo,
-      BenchmarkCode *Result)
-      : MCStreamer(*Context), RegNameToRegNo(RegNameToRegNo), Result(Result) {}
+  explicit BenchmarkCodeStreamer(MCContext *Context, const LLVMState &State,
+                                 BenchmarkCode *Result)
+      : MCStreamer(*Context), State(State), Result(Result) {}
 
   // Implementation of the MCStreamer interface. We only care about
   // instructions.
@@ -207,15 +208,17 @@ class BenchmarkCodeStreamer : public MCStreamer, public AsmCommentConsumer {
                     Align ByteAlignment, SMLoc Loc) override {}
 
   unsigned findRegisterByName(const StringRef RegName) const {
-    auto Iter = RegNameToRegNo.find(RegName);
-    if (Iter != RegNameToRegNo.end())
-      return Iter->second;
-    errs() << "'" << RegName
-           << "' is not a valid register name for the target\n";
-    return 0;
+    std::optional<MCRegister> RegisterNumber =
+        State.getRegisterNumberFromName(RegName);
+    if (!RegisterNumber.has_value()) {
+      errs() << "'" << RegName
+             << "' is not a valid register name for the target\n";
+      return MCRegister::NoRegister;
+    }
+    return *RegisterNumber;
   }
 
-  const DenseMap<StringRef, MCRegister> &RegNameToRegNo;
+  const LLVMState &State;
   BenchmarkCode *const Result;
   unsigned InvalidComments = 0;
 };
@@ -248,8 +251,7 @@ Expected<std::vector<BenchmarkCode>> readSnippets(const LLVMState &State,
       TM.getTarget().createMCObjectFileInfo(Context, /*PIC=*/false));
   Context.setObjectFileInfo(ObjectFileInfo.get());
   Context.initInlineSourceManager();
-  BenchmarkCodeStreamer Streamer(&Context, State.getRegNameToRegNoMapping(),
-                                 &Result);
+  BenchmarkCodeStreamer Streamer(&Context, State, &Result);
 
   std::string Error;
   raw_string_ostream ErrorStream(Error);