[Bolt] Solving pie support issue

bolt/lib/Core/BinaryContext.cpp

@@ -1924,10 +1924,27 @@ BinaryContext::getBaseAddressForMapping(uint64_t MMapAddress,
   // Find a segment with a matching file offset.
   for (auto &KV : SegmentMapInfo) {
     const SegmentInfo &SegInfo = KV.second;
-    if (alignDown(SegInfo.FileOffset, SegInfo.Alignment) == FileOffset) {
-      // Use segment's aligned memory offset to calculate the base address.
-      const uint64_t MemOffset = alignDown(SegInfo.Address, SegInfo.Alignment);
-      return MMapAddress - MemOffset;
+    // FileOffset is got from perf event,
+    // and it is equal to alignDown(SegInfo.FileOffset, pagesize).
+    // If the pagesize is not equal to SegInfo.Alignment.
+    // FileOffset and SegInfo.FileOffset should be aligned first,
+    // and then judge whether they are equal.
+    if (alignDown(SegInfo.FileOffset, SegInfo.Alignment) ==
+        alignDown(FileOffset, SegInfo.Alignment)) {
+      // The function's offset from base address in VAS is aligned by pagesize
+      // instead of SegInfo.Alignment. Pagesize can't be got from perf events.
+      // However, The ELF document says that SegInfo.FileOffset should equal
+      // to SegInfo.Address, modulo the pagesize.
+      // Reference: https://refspecs.linuxfoundation.org/elf/elf.pdf
+
+      // So alignDown(SegInfo.Address, pagesize) can be calculated by:
+      // alignDown(SegInfo.Address, pagesize)
+      //   = SegInfo.Address - (SegInfo.Address % pagesize)
+      //   = SegInfo.Address - (SegInfo.FileOffset % pagesize)
+      //   = SegInfo.Address - SegInfo.FileOffset +
+      //     alignDown(SegInfo.FileOffset, pagesize)
+      //   = SegInfo.Address - SegInfo.FileOffset + FileOffset
+      return MMapAddress - (SegInfo.Address - SegInfo.FileOffset + FileOffset);
     }
   }
 

bolt/test/perf2bolt/Inputs/perf_test.c

@@ -0,0 +1,26 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+int add(int a, int b) { return a + b; }
+int minus(int a, int b) { return a - b; }
+int multiple(int a, int b) { return a * b; }
+int divide(int a, int b) {
+  if (b == 0)
+    return 0;
+  return a / b;
+}
+
+int main() {
+  int a = 16;
+  int b = 8;
+
+  for (int i = 1; i < 100000; i++) {
+    add(a, b);
+    minus(a, b);
+    multiple(a, b);
+    divide(a, b);
+  }
+
+  return 0;
+}

bolt/test/perf2bolt/Inputs/perf_test.lds

@@ -0,0 +1,13 @@
+SECTIONS {
+  . = SIZEOF_HEADERS;
+  .interp : { *(.interp) }
+  .note.gnu.build-id : { *(.note.gnu.build-id) }
+  . = 0x212e8;
+  .dynsym         : { *(.dynsym) }
+  . = 0x31860;
+  .text : { *(.text*) }
+  . = 0x41c20;
+  .fini_array : { *(.fini_array) }
+  . = 0x54e18;
+  .data : { *(.data) }
+}

bolt/test/perf2bolt/lit.local.cfg

@@ -0,0 +1,4 @@
+import shutil
+
+if shutil.which("perf") != None:
+    config.available_features.add("perf")

bolt/test/perf2bolt/perf_test.test

@@ -0,0 +1,17 @@
+# Check perf2bolt binary function which was compiled with pie
+
+REQUIRES: system-linux, perf
+
+RUN: %clang %S/Inputs/perf_test.c -fuse-ld=lld -Wl,--script=%S/Inputs/perf_test.lds -o %t
+RUN: perf record -e cycles:u -o %t2 -- %t
+RUN: perf2bolt %t -p=%t2 -o %t3 -nl -ignore-build-id 2>&1 | FileCheck %s
+
+CHECK-NOT: PERF2BOLT-ERROR
+CHECK-NOT: !! WARNING !! This high mismatch ratio indicates the input binary is probably not the same binary used during profiling collection.
+
+RUN: %clang %S/Inputs/perf_test.c -no-pie -fuse-ld=lld -o %t4
+RUN: perf record -e cycles:u -o %t5 -- %t4
+RUN: perf2bolt %t4 -p=%t5 -o %t6 -nl -ignore-build-id 2>&1 | FileCheck %s --check-prefix=CHECK-NO-PIE
+
+CHECK-NO-PIE-NOT: PERF2BOLT-ERROR
+CHECK-NO-PIE-NOT: !! WARNING !! This high mismatch ratio indicates the input binary is probably not the same binary used during profiling collection.

bolt/unittests/Core/BinaryContext.cpp

@@ -123,3 +123,24 @@ TEST_P(BinaryContextTester, BaseAddress) {
   BaseAddress = BC->getBaseAddressForMapping(0x7f13f5556000, 0x137a000);
   ASSERT_FALSE(BaseAddress.has_value());
 }
+
+TEST_P(BinaryContextTester, BaseAddress2) {
+  // Check that base address calculation is correct for a binary if the
+  // alignment in ELF file are different from pagesize.
+  // The segment layout is as follows:
+  BC->SegmentMapInfo[0] = SegmentInfo{0, 0x2177c, 0, 0x2177c, 0x10000};
+  BC->SegmentMapInfo[0x31860] =
+      SegmentInfo{0x31860, 0x370, 0x21860, 0x370, 0x10000};
+  BC->SegmentMapInfo[0x41c20] =
+      SegmentInfo{0x41c20, 0x1f8, 0x21c20, 0x1f8, 0x10000};
+  BC->SegmentMapInfo[0x54e18] =
+      SegmentInfo{0x54e18, 0x51, 0x24e18, 0x51, 0x10000};
+
+  std::optional<uint64_t> BaseAddress =
+      BC->getBaseAddressForMapping(0xaaaaea444000, 0x21000);
+  ASSERT_TRUE(BaseAddress.has_value());
+  ASSERT_EQ(*BaseAddress, 0xaaaaea413000ULL);
+
+  BaseAddress = BC->getBaseAddressForMapping(0xaaaaea444000, 0x11000);
+  ASSERT_FALSE(BaseAddress.has_value());
+}