[llvm] [llvm-exegesis] Add support for memory annotations in yaml (PR #76665)

[Aiden Grossman via llvm-commits]

https://github.com/boomanaiden154 created https://github.com/llvm/llvm-project/pull/76665

This patch adds support for exporting memory annotations in the yaml. This isn't necessary for analysis mode, but it can be useful to have information about the memory setup for the snippet execution in the yaml output. This also makes the output for memory annotations consistent with other annotations.

>From 910cf590bbae3645fb0fcbafb0adfbcdab14c400 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <agrossman154 at yahoo.com>
Date: Sun, 31 Dec 2023 18:00:32 -0800
Subject: [PATCH] [llvm-exegesis] Add support for memory annotations in yaml

This patch adds support for exporting memory annotations in the yaml.
This isn't necessary for analysis mode, but it can be useful to have
information about the memory setup for the snippet execution in the yaml
output. This also makes the output for memory annotations consistent
with other annotations.
---
 .../llvm-exegesis/lib/BenchmarkResult.cpp     | 60 +++++++++++++++++++
 .../llvm-exegesis/X86/BenchmarkResultTest.cpp |  2 +
 2 files changed, 62 insertions(+)

diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp
index 02c4da11e032d6..cc1edd67172fe8 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.cpp
@@ -258,6 +258,64 @@ template <> struct ScalarTraits<exegesis::RegisterValue> {
   static const bool flow = true;
 };
 
+template <> struct ScalarTraits<APInt> {
+  static constexpr const unsigned kRadix = 16;
+  static constexpr const bool kSigned = false;
+
+  static void output(const APInt &Input, void *Ctx, raw_ostream &Out) {
+    std::string OutputString = toString(Input, kRadix, kSigned);
+    // The snippet parsing infrastructure and other frontends should only
+    // create APInts with a bit width of a multiple of four for memory
+    // values.
+    assert(Input.getBitWidth() % 4 == 0);
+    for (size_t I = OutputString.size(); I < Input.getBitWidth() / 4; ++I) {
+      Out << "0";
+    }
+    Out << OutputString;
+  }
+
+  static StringRef input(StringRef String, void *Ctx, APInt &Output) {
+    Output = APInt(String.size() * 4, String, kRadix);
+    return StringRef();
+  }
+
+  static QuotingType mustQuote(StringRef) { return QuotingType::None; }
+};
+
+template <> struct MappingTraits<exegesis::MemoryValue> {
+  static void mapping(IO &Io, exegesis::MemoryValue &MV) {
+    Io.mapRequired("size", MV.SizeBytes);
+    Io.mapRequired("value", MV.Value);
+  }
+};
+
+template <>
+struct CustomMappingTraits<
+    std::unordered_map<std::string, exegesis::MemoryValue>> {
+  static void
+  inputOne(IO &Io, StringRef Key,
+           std::unordered_map<std::string, exegesis::MemoryValue> &MVM) {
+    Io.mapRequired(Key.str().c_str(), MVM[Key.str()]);
+  }
+
+  static void
+  output(IO &Io, std::unordered_map<std::string, exegesis::MemoryValue> &MVM) {
+    for (auto &MV : MVM)
+      Io.mapRequired(MV.first.c_str(), MV.second);
+  }
+};
+
+template <> struct SequenceElementTraits<exegesis::MemoryMapping> {
+  static const bool flow = false;
+};
+
+template <> struct MappingTraits<exegesis::MemoryMapping> {
+  static void mapping(IO &Io, exegesis::MemoryMapping &MM) {
+    Io.mapRequired("address", MM.Address);
+    Io.mapRequired("memory_value", MM.MemoryValueName);
+  }
+};
+
 template <>
 struct MappingContextTraits<exegesis::BenchmarkKey, YamlContext> {
   static void mapping(IO &Io, exegesis::BenchmarkKey &Obj,
@@ -266,6 +324,8 @@ struct MappingContextTraits<exegesis::BenchmarkKey, YamlContext> {
     Io.mapRequired("instructions", Obj.Instructions);
     Io.mapOptional("config", Obj.Config);
     Io.mapRequired("register_initial_values", Obj.RegisterInitialValues);
+    Io.mapRequired("memory_values", Obj.MemoryValues);
+    Io.mapRequired("memory_mappings", Obj.MemoryMappings);
   }
 };
 
diff --git a/llvm/unittests/tools/llvm-exegesis/X86/BenchmarkResultTest.cpp b/llvm/unittests/tools/llvm-exegesis/X86/BenchmarkResultTest.cpp
index 6c558b59be982d..4671aee9da8592 100644
--- a/llvm/unittests/tools/llvm-exegesis/X86/BenchmarkResultTest.cpp
+++ b/llvm/unittests/tools/llvm-exegesis/X86/BenchmarkResultTest.cpp
@@ -71,6 +71,8 @@ TEST(BenchmarkResultTest, WriteToAndReadFromDisk) {
   ToDisk.Key.RegisterInitialValues = {
       RegisterValue{X86::AL, APInt(8, "-1", 10)},
       RegisterValue{X86::AH, APInt(8, "123", 10)}};
+  ToDisk.Key.MemoryValues = {{"test1", {APInt(64, 0xff), 4096, 0}}};
+  ToDisk.Key.MemoryMappings = {{0x4000, "test1"}};
   ToDisk.Mode = Benchmark::Latency;
   ToDisk.CpuName = "cpu_name";
   ToDisk.LLVMTriple = "llvm_triple";