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

Before this patch, in subprocess mode, llvm-exegesis setup the stack
pointer register with the rest of the registers when it was requested by
the user. This would cause a segfault when the instructions to start the
perf counter ran as they use the stack to preserve the three registers
needed to make the syscall. This patch moves the setup of the stack
register to after the configuration of the perf counter to fix this
issue so that we have a valid stack pointer for all the preceeding
operations.

Regression test added.

This fixes #72193.

Author: boomanaiden154

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/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;
 }