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

[Aiden Grossman via llvm-commits]

Author: Aiden Grossman
Date: 2023-11-30T15:27:26-08:00
New Revision: 668865789620f390fbad4d7093ed8ca6eb932c31

URL: https://github.com/llvm/llvm-project/commit/668865789620f390fbad4d7093ed8ca6eb932c31
DIFF: https://github.com/llvm/llvm-project/commit/668865789620f390fbad4d7093ed8ca6eb932c31.diff

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

This reverts commit 52b4b357268ac114f7972e37da1e5954f322df09.

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.

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

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

Removed: 
    


################################################################################
diff  --git a/llvm/test/tools/llvm-exegesis/X86/latency/subprocess-rsp.s b/llvm/test/tools/llvm-exegesis/X86/latency/subprocess-rsp.s
new file mode 100644
index 000000000000000..cc23cdee3296103
--- /dev/null
+++ b/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'

diff  --git a/llvm/tools/llvm-exegesis/lib/Assembler.cpp b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
index e64f4e64ebb4516..e17d239faa47e21 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.cpp
+++ b/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;
 }