[llvm-exegesis] Save target state before running the benchmark.

Some benchmarked instructions might set target state. Preserve this
state. See PR26418.

Differential Revision: https://reviews.llvm.org/D90592

llvm/test/tools/llvm-exegesis/X86/uops-FLDENVm.s

@@ -0,0 +1,6 @@
+# RUN: llvm-exegesis -mode=uops -opcode-name=FLDENVm,FLDL2E -repetition-mode=duplicate | FileCheck %s
+
+CHECK:      mode:            uops
+CHECK-NEXT: key:
+CHECK-NEXT:   instructions:
+CHECK-NEXT:     FLDENVm

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

@@ -71,10 +71,10 @@ class FunctionExecutorImpl : public BenchmarkRunner::FunctionExecutor {
     SmallVector<StringRef, 2> CounterNames;
     StringRef(Counters).split(CounterNames, '+');
     char *const ScratchPtr = Scratch->ptr();
+    const ExegesisTarget &ET = State.getExegesisTarget();
     for (auto &CounterName : CounterNames) {
       CounterName = CounterName.trim();
-      auto CounterOrError =
-          State.getExegesisTarget().createCounter(CounterName, State);
+      auto CounterOrError = ET.createCounter(CounterName, State);
 
       if (!CounterOrError)
         return CounterOrError.takeError();
@@ -93,6 +93,7 @@ class FunctionExecutorImpl : public BenchmarkRunner::FunctionExecutor {
                 .concat(std::to_string(Reserved)));
       Scratch->clear();
       {
+        auto PS = ET.withSavedState();
         CrashRecoveryContext CRC;
         CrashRecoveryContext::Enable();
         const bool Crashed = !CRC.RunSafely([this, Counter, ScratchPtr]() {
@@ -101,6 +102,7 @@ class FunctionExecutorImpl : public BenchmarkRunner::FunctionExecutor {
           Counter->stop();
         });
         CrashRecoveryContext::Disable();
+        PS.reset();
         if (Crashed) {
           std::string Msg = "snippet crashed while running";
 #ifdef LLVM_ON_UNIX

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

@@ -147,6 +147,8 @@ const PfmCountersInfo &ExegesisTarget::getPfmCounters(StringRef CpuName) const {
   return *Found->PCI;
 }
 
+ExegesisTarget::SavedState::~SavedState() {} // anchor.
+
 namespace {
 
 // Default implementation.

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

@@ -172,6 +172,16 @@ class ExegesisTarget {
   // counters are defined for this CPU).
   const PfmCountersInfo &getPfmCounters(StringRef CpuName) const;
 
+  // Saves the CPU state that needs to be preserved when running a benchmark,
+  // and returns and RAII object that restores the state on destruction.
+  // By default no state is preserved.
+  struct SavedState {
+    virtual ~SavedState();
+  };
+  virtual std::unique_ptr<SavedState> withSavedState() const {
+    return std::make_unique<SavedState>();
+  }
+
 private:
   virtual bool matchesArch(Triple::ArchType Arch) const = 0;
 

llvm/tools/llvm-exegesis/lib/X86/Target.cpp

@@ -23,6 +23,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 
+#include <immintrin.h>
 #include <memory>
 #include <string>
 #include <vector>
@@ -594,6 +595,25 @@ void ConstantInliner::initStack(unsigned Bytes) {
 
 namespace {
 
+class X86SavedState : public ExegesisTarget::SavedState {
+public:
+  X86SavedState() { _fxsave64(FPState); }
+
+  ~X86SavedState() {
+    // Restoring the X87 state does not flush pending exceptions, make sure
+    // these exceptions are flushed now.
+#if defined(_MSC_VER)
+    _clearfp();
+#elif defined(__GNUC__)
+    asm volatile("fwait");
+#endif
+    _fxrstor64(FPState);
+  }
+
+private:
+  alignas(16) char FPState[512];
+};
+
 class ExegesisX86Target : public ExegesisTarget {
 public:
   ExegesisX86Target() : ExegesisTarget(X86CpuPfmCounters) {}
@@ -691,6 +711,10 @@ class ExegesisX86Target : public ExegesisTarget {
 #endif
   }
 
+  std::unique_ptr<SavedState> withSavedState() const override {
+    return std::make_unique<X86SavedState>();
+  }
+
   static const unsigned kUnavailableRegisters[4];
 };