[llvm-exegesis] Add in snippet address annotation (llvm#74218)

Author: boomanaiden154

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
 

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

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'

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 {

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

@@ -416,9 +416,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);
 

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

@@ -131,6 +131,19 @@ 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>
+      if (!to_integer<intptr_t>(CommentText.trim(), Result->Key.SnippetAddress,
+                                16)) {
+        errs() << "invalid comment 'LLVM-EXEGESIS-SNIPPET-ADDRESS "
+               << CommentText
+               << "', expected <ADDRESS> to contain a valid integer in "
+                  "hexadecimal format";
+        ++InvalidComments;
+        return;
+      }
+      return;
+    }
   }
 
   unsigned numInvalidComments() const { return InvalidComments; }

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");
     }
   }

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