[llvm-exegesis] Set stack pointer register after starting perf counter

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

@@ -0,0 +1,18 @@
+# REQUIRES: exegesis-can-execute-x86_64
+
+# 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/SectionMemoryManager.h"
@@ -59,7 +60,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())
@@ -70,6 +81,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;
 }