[llvm-exegesis] Add in snippet address annotation

[boomanaiden154]

This patch adds in a LLVM-EXEGESIS-SNIPPET-ADDRESS annotation that can be used in snippets executed by llvm-exegesis. This moves over the code that exegesis runs (including the setup) to start at the address specified. This allows for the execution of snippets that perform actions like RIP-relative variable referencing.

[llvmbot]

@llvm/pr-subscribers-tools-llvm-exegesis

Author: Aiden Grossman (boomanaiden154)

Changes

This patch adds in a LLVM-EXEGESIS-SNIPPET-ADDRESS annotation that can be used in snippets executed by llvm-exegesis. This moves over the code that exegesis runs (including the setup) to start at the address specified. This allows for the execution of snippets that perform actions like RIP-relative variable referencing.

---

Full diff: https://github.com/llvm/llvm-project/pull/74218.diff

8 Files Affected:

• (modified) llvm/test/tools/llvm-exegesis/X86/latency/memory-annotations-unsupported.s (+1-1)
• (added) llvm/test/tools/llvm-exegesis/X86/latency/snippet-address-annotations-unsupported.s (+9)
• (added) llvm/test/tools/llvm-exegesis/X86/latency/subprocess-address-annotation.s (+26)
• (modified) llvm/tools/llvm-exegesis/lib/BenchmarkResult.h (+3)
• (modified) llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp (+10-2)
• (modified) llvm/tools/llvm-exegesis/lib/SnippetFile.cpp (+6)
• (modified) llvm/tools/llvm-exegesis/llvm-exegesis.cpp (+4-2)
• (modified) llvm/unittests/tools/llvm-exegesis/X86/SnippetFileTest.cpp (+10)

diff --git a/llvm/test/tools/llvm-exegesis/X86/latency/memory-annotations-unsupported.s b/llvm/test/tools/llvm-exegesis/X86/latency/memory-annotations-unsupported.s
index 5d06a0ceff5c7..c483d28adf0e1 100644
--- a/llvm/test/tools/llvm-exegesis/X86/latency/memory-annotations-unsupported.s
+++ b/llvm/test/tools/llvm-exegesis/X86/latency/memory-annotations-unsupported.s
@@ -2,7 +2,7 @@
 
 # RUN: not llvm-exegesis -mtriple=x86_64-unknown-unknown -snippets-file=%s -mode=latency 2>&1 | FileCheck %s
 
-# CHECK: llvm-exegesis error: Memory annotations are only supported in subprocess execution mode
+# CHECK: llvm-exegesis error: Memory and snippet address annotations are only supported in subprocess execution mode
 
 # LLVM-EXEGESIS-MEM-DEF test1 4096 ff
 
diff --git a/llvm/test/tools/llvm-exegesis/X86/latency/snippet-address-annotations-unsupported.s b/llvm/test/tools/llvm-exegesis/X86/latency/snippet-address-annotations-unsupported.s
new file mode 100644
index 0000000000000..c38c41d2aed9b
--- /dev/null
+++ b/llvm/test/tools/llvm-exegesis/X86/latency/snippet-address-annotations-unsupported.s
@@ -0,0 +1,9 @@
+# REQUIRES: exegesis-can-measure-latency, x86_64-linux
+
+# RUN: not llvm-exegesis -mtriple=x86_64-unknown-unknown -snippets-file=%s -mode=latency 2>&1 | FileCheck %s
+
+# CHECK: llvm-exegesis error: Memory and snippet address annotations are only supported in subprocess execution mode
+
+# LLVM-EXEGESIS-SNIPPET-ADDRESS 10000
+
+movq $0, %rax
diff --git a/llvm/test/tools/llvm-exegesis/X86/latency/subprocess-address-annotation.s b/llvm/test/tools/llvm-exegesis/X86/latency/subprocess-address-annotation.s
new file mode 100644
index 0000000000000..1b104b30bb288
--- /dev/null
+++ b/llvm/test/tools/llvm-exegesis/X86/latency/subprocess-address-annotation.s
@@ -0,0 +1,26 @@
+# 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 the code is loaded in at the expected address.
+
+# LLVM-EXEGESIS-SNIPPET-ADDRESS 20000
+# LLVM-EXEGESIS-DEFREG RAX 0
+# LLVM-EXEGESIS-DEFREG R14 127
+# LLVM-EXEGESIS-DEFREG R15 0
+# LLVM-EXEGESIS-DEFREG RDI 0
+
+# Load the instruction pointer and round down to the nearest page as there
+# will be some setup code loaded in before this part begins to execute.
+lea 0(%rip), %rax
+shrq $12, %rax
+shlq $12, %rax
+
+cmpq $0x20000, %rax
+cmovneq %r14, %r15
+
+movq $60, %rax
+movq %r15, %rdi
+syscall
+
+# CHECK-NOT: error:           'Child benchmarking process exited with non-zero exit code: Child process returned with unknown exit code'
diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h
index 38111519a2c89..0d08febae20cb 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h
@@ -70,6 +70,9 @@ struct BenchmarkKey {
   // An opaque configuration, that can be used to separate several benchmarks of
   // the same instruction under different configurations.
   std::string Config;
+  // The address that the snippet should be loaded in at if the execution mode
+  // being used supports it.
+  intptr_t SnippetAddress = 0;
 };
 
 struct BenchmarkMeasure {
diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp b/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp
index 85375dec2a44c..700f9beb962ec 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp
@@ -404,9 +404,17 @@ class SubProcessFunctionExecutorImpl
 #endif // GLIBC_INITS_RSEQ
 
     size_t FunctionDataCopySize = this->Function.FunctionBytes.size();
+    void *MapAddress = NULL;
+    int MapFlags = MAP_PRIVATE | MAP_ANONYMOUS;
+
+    if (Key.SnippetAddress != 0) {
+      MapAddress = reinterpret_cast<void *>(Key.SnippetAddress);
+      MapFlags |= MAP_FIXED_NOREPLACE;
+    }
+
     char *FunctionDataCopy =
-        (char *)mmap(NULL, FunctionDataCopySize, PROT_READ | PROT_WRITE,
-                     MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
+        (char *)mmap(MapAddress, FunctionDataCopySize, PROT_READ | PROT_WRITE,
+                     MapFlags, 0, 0);
     if ((intptr_t)FunctionDataCopy == -1)
       exit(ChildProcessExitCodeE::FunctionDataMappingFailed);
 
diff --git a/llvm/tools/llvm-exegesis/lib/SnippetFile.cpp b/llvm/tools/llvm-exegesis/lib/SnippetFile.cpp
index d85a9f190655a..a140f7b2ba77b 100644
--- a/llvm/tools/llvm-exegesis/lib/SnippetFile.cpp
+++ b/llvm/tools/llvm-exegesis/lib/SnippetFile.cpp
@@ -131,6 +131,12 @@ class BenchmarkCodeStreamer : public MCStreamer, public AsmCommentConsumer {
       Result->Key.MemoryMappings.push_back(std::move(MemMap));
       return;
     }
+    if (CommentText.consume_front("SNIPPET-ADDRESS")) {
+      // LLVM-EXEGESIS-SNIPPET-ADDRESS <address>
+      Result->Key.SnippetAddress =
+          std::stol(CommentText.trim().str(), nullptr, 16);
+      return;
+    }
   }
 
   unsigned numInvalidComments() const { return InvalidComments; }
diff --git a/llvm/tools/llvm-exegesis/llvm-exegesis.cpp b/llvm/tools/llvm-exegesis/llvm-exegesis.cpp
index 261335a817d06..66bc94271756c 100644
--- a/llvm/tools/llvm-exegesis/llvm-exegesis.cpp
+++ b/llvm/tools/llvm-exegesis/llvm-exegesis.cpp
@@ -533,8 +533,10 @@ void benchmarkMain() {
     for (const auto &Configuration : Configurations) {
       if (ExecutionMode != BenchmarkRunner::ExecutionModeE::SubProcess &&
           (Configuration.Key.MemoryMappings.size() != 0 ||
-           Configuration.Key.MemoryValues.size() != 0))
-        ExitWithError("Memory annotations are only supported in subprocess "
+           Configuration.Key.MemoryValues.size() != 0 ||
+           Configuration.Key.SnippetAddress != 0))
+        ExitWithError("Memory and snippet address annotations are only "
+                      "supported in subprocess "
                       "execution mode");
     }
   }
diff --git a/llvm/unittests/tools/llvm-exegesis/X86/SnippetFileTest.cpp b/llvm/unittests/tools/llvm-exegesis/X86/SnippetFileTest.cpp
index 2bab3aa15baad..52eb211c37272 100644
--- a/llvm/unittests/tools/llvm-exegesis/X86/SnippetFileTest.cpp
+++ b/llvm/unittests/tools/llvm-exegesis/X86/SnippetFileTest.cpp
@@ -209,6 +209,16 @@ TEST_F(X86SnippetFileTest, MemoryMappingNoDefinition) {
   consumeError(std::move(Error));
 }
 
+TEST_F(X86SnippetFileTest, SnippetAddress) {
+  auto Snippets = TestCommon(R"(
+    # LLVM-EXEGESIS-SNIPPET-ADDRESS 0x10000
+  )");
+  ASSERT_TRUE(static_cast<bool>(Snippets));
+  EXPECT_THAT(*Snippets, SizeIs(1));
+  const auto &Snippet = (*Snippets)[0];
+  EXPECT_EQ(Snippet.Key.SnippetAddress, 0x10000);
+}
+
 } // namespace
 } // namespace exegesis
 } // namespace llvm