fixup! Address review comments

Author: mshockwave

llvm/include/llvm/Support/Compression.h

@@ -15,6 +15,7 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/YAMLTraits.h"
 
 namespace llvm {
 template <typename T> class SmallVectorImpl;
@@ -126,6 +127,16 @@ Error decompress(DebugCompressionType T, ArrayRef<uint8_t> Input,
 
 } // End of namespace compression
 
+namespace yaml {
+// Related YAML traits.
+template <> struct ScalarEnumerationTraits<compression::Format> {
+  static void enumeration(IO &Io, compression::Format &Format) {
+    Io.enumCase(Format, "zstd", compression::Format::Zstd);
+    Io.enumCase(Format, "zlib", compression::Format::Zlib);
+  }
+};
+} // namespace yaml
+
 } // End of namespace llvm
 
 #endif

llvm/test/tools/llvm-exegesis/serialize-obj-file.test

@@ -1,16 +1,22 @@
-# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --benchmark-phase=assemble-measured-code --mode=latency --benchmarks-file=%t.yaml
+# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --exegesis-serialize-benchmarks --benchmark-phase=assemble-measured-code \
+# RUN:    --mode=latency --benchmarks-file=%t.yaml
 # RUN: FileCheck --input-file=%t.yaml %s --check-prefixes=CHECK,SERIALIZE
 # RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --run-measurement=%t.yaml --mode=latency --benchmark-phase=dry-run-measurement --use-dummy-perf-counters \
 # RUN:    --dump-object-to-disk=%t.o | FileCheck %s --check-prefixes=CHECK,DESERIALIZE
 # RUN: llvm-objdump -d %t.o | FileCheck %s --check-prefix=OBJDUMP
-# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --benchmark-phase=dry-run-measurement --use-dummy-perf-counters | \
-# RUN:    FileCheck %s --check-prefix=NO-SERIALIZE
-# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --benchmark-phase=assemble-measured-code --repetition-mode=min | \
-# RUN:    FileCheck %s --check-prefix=NO-SERIALIZE
-# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --benchmark-phase=assemble-measured-code --repetition-mode=middle-half-loop | \
-# RUN:    FileCheck %s --check-prefix=NO-SERIALIZE
-# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --benchmark-phase=assemble-measured-code --repetition-mode=middle-half-duplicate | \
+
+# We should not serialie benchmarks by default.
+# RUN: llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --benchmark-phase=assemble-measured-code --mode=latency | \
 # RUN:    FileCheck %s --check-prefix=NO-SERIALIZE
+
+# We currently don't support serialization for repetition modes that require more than one snippets.
+# RUN: not llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --exegesis-serialize-benchmarks --benchmark-phase=assemble-measured-code \
+# RUN:    --repetition-mode=min 2>&1 | FileCheck %s --check-prefix=NOT-SUPPORTED
+# RUN: not llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --exegesis-serialize-benchmarks --benchmark-phase=assemble-measured-code \
+# RUN:    --repetition-mode=middle-half-loop 2>&1 | FileCheck %s --check-prefix=NOT-SUPPORTED
+# RUN: not llvm-exegesis -mtriple=riscv64 -mcpu=sifive-p470 --opcode-name=SH3ADD --mode=latency --exegesis-serialize-benchmarks --benchmark-phase=assemble-measured-code \
+# RUN:    --repetition-mode=middle-half-duplicate 2>&1 | FileCheck %s --check-prefix=NOT-SUPPORTED
+
 # REQUIRES: riscv-registered-target && native-registered-exegesis-target
 # REQUIRES: zlib || zstd
 
@@ -29,6 +35,7 @@
 
 # OBJDUMP: sh3add
 
-# Negative tests: we shouldn't serialize object files in some scenarios.
+# Negative tests.
 
+# NOT-SUPPORTED: -exegesis-serialize-benchmarks currently only supports -repetition-mode of loop and duplicate.
 # NO-SERIALIZE-NOT: object_file:

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

@@ -287,13 +287,6 @@ template <> struct ScalarTraits<exegesis::RegisterValue> {
   static const bool flow = true;
 };
 
-template <> struct ScalarEnumerationTraits<compression::Format> {
-  static void enumeration(IO &Io, compression::Format &Format) {
-    Io.enumCase(Format, "zstd", compression::Format::Zstd);
-    Io.enumCase(Format, "zlib", compression::Format::Zlib);
-  }
-};
-
 template <> struct MappingContextTraits<exegesis::BenchmarkKey, YamlContext> {
   static void mapping(IO &Io, exegesis::BenchmarkKey &Obj,
                       YamlContext &Context) {

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

@@ -19,6 +19,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Config/llvm-config.h" // for LLVM_ON_UNIX
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
@@ -51,6 +52,14 @@
 #endif // __linux__
 
 namespace llvm {
+
+static cl::opt<bool>
+    SerializeBenchmarks("exegesis-serialize-benchmarks",
+                        cl::desc("Generate fully-serialized benchmarks "
+                                 "that can later be deserialized and "
+                                 "resuming the measurement."),
+                        cl::init(false));
+
 namespace exegesis {
 
 BenchmarkRunner::BenchmarkRunner(const LLVMState &State, Benchmark::ModeE Mode,
@@ -664,16 +673,19 @@ BenchmarkRunner::getRunnableConfiguration(
                         LoopBodySize, GenerateMemoryInstructions);
     if (Error E = Snippet.takeError())
       return std::move(E);
-    // There is no need to serialize/deserialize the object file if we're
-    // simply running end-to-end measurements.
-    // Same goes for any repetition mode that requires more than a single
-    // snippet.
-    if (BenchmarkPhaseSelector < BenchmarkPhaseSelectorE::Measure &&
-        (RepetitionMode == Benchmark::Loop ||
-         RepetitionMode == Benchmark::Duplicate)) {
+
+    // Generate fully-serialized benchmarks.
+    if (SerializeBenchmarks) {
+      if (RepetitionMode != Benchmark::Loop &&
+          RepetitionMode != Benchmark::Duplicate)
+        return make_error<Failure>(
+            "-exegesis-serialize-benchmarks currently "
+            "only supports -repetition-mode of loop and duplicate.");
+
       if (Error E = BenchmarkResult.setObjectFile(*Snippet))
         return std::move(E);
     }
+
     RC.ObjectFile = getObjectFromBuffer(*Snippet);
   }