Reland "[llvm-exegesis] Set stack pointer register after starting perf counter (#72489)"

This reverts commit 52b4b35726.

This commit was causing build failures on bots without libpfm installed
as the CHECK requires were not set correctly. The test added should only
run on machines with libpfm enabled due to the limitations currently
imposed by the subprocess execution mode.

llvm/test/tools/llvm-exegesis/X86/latency/subprocess-rsp.s

@@ -0,0 +1,18 @@
+# REQUIRES: exegesis-can-measure-latency, x86_64-linux
+
+# RUN: llvm-exegesis -mtriple=x86_64-unknown-unknown -mode=latency -snippets-file=%s -execution-mode=subprocess | FileCheck %s
+
+# Check that we can set the value of RSP in subprocess mode without
+# segfaulting as we need to restore it after the rest of the setup is
+# complete to prevent loading from the stack where we set it instead
+# of where the stack actuall is.
+
+# LLVM-EXEGESIS-MEM-DEF test1 4096 2147483647
+# LLVM-EXEGESIS-MEM-MAP test1 1048576
+# LLVM-EXEGESIS-DEFREG RAX 100000
+# LLVM-EXEGESIS-DEFREG R14 100000
+# LLVM-EXEGESIS-DEFREG RSP 100000
+
+movq %r14, (%rax)
+
+# CHECK-NOT: error:           'The benchmarking subprocess sent unexpected signal: Segmentation fault'

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

@@ -19,6 +19,7 @@
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
@@ -60,7 +61,17 @@ static bool generateSnippetSetupCode(
     }
     BBF.addInstructions(ET.setStackRegisterToAuxMem());
   }
+  Register StackPointerRegister = BBF.MF.getSubtarget()
+                                      .getTargetLowering()
+                                      ->getStackPointerRegisterToSaveRestore();
   for (const RegisterValue &RV : RegisterInitialValues) {
+    if (GenerateMemoryInstructions) {
+      // If we're generating memory instructions, don't load in the value for
+      // the register with the stack pointer as it will be used later to finish
+      // the setup.
+      if (RV.Register == StackPointerRegister)
+        continue;
+    }
     // Load a constant in the register.
     const auto SetRegisterCode = ET.setRegTo(*MSI, RV.Register, RV.Value);
     if (SetRegisterCode.empty())
@@ -71,6 +82,18 @@ static bool generateSnippetSetupCode(
 #ifdef HAVE_LIBPFM
     BBF.addInstructions(ET.configurePerfCounter(PERF_EVENT_IOC_RESET, true));
 #endif // HAVE_LIBPFM
+    for (const RegisterValue &RV : RegisterInitialValues) {
+      // Load in the stack register now as we're done using it elsewhere
+      // and need to set the value in preparation for executing the
+      // snippet.
+      if (RV.Register == StackPointerRegister)
+        continue;
+      const auto SetRegisterCode = ET.setRegTo(*MSI, RV.Register, RV.Value);
+      if (SetRegisterCode.empty())
+        IsSnippetSetupComplete = false;
+      BBF.addInstructions(SetRegisterCode);
+      break;
+    }
   }
   return IsSnippetSetupComplete;
 }