[BOLT] Align constant islands to 8 bytes

AArch64 requires CI to be aligned to 8 bytes due to access instructions
restrictions. E.g. the ldr with imm, where imm must be aligned to 8 bytes.

Differential Revision: https://reviews.llvm.org/D122065

5 files changed, 63 insertions, 4 deletions

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 3d9b99496a63..99674ce5c5bf 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -172,6 +172,9 @@ public:
 
     mutable MCSymbol *FunctionConstantIslandLabel{nullptr};
     mutable MCSymbol *FunctionColdConstantIslandLabel{nullptr};
+
+    // Returns constant island alignment
+    uint16_t getAlignment() const { return sizeof(uint64_t); }
   };
 
   static constexpr uint64_t COUNT_NO_PROFILE =
@@ -2047,6 +2050,10 @@ public:
     return *std::prev(CodeIter) <= *DataIter;
   }
 
+  uint16_t getConstantIslandAlignment() const {
+    return Islands ? Islands->getAlignment() : 1;
+  }
+
   uint64_t
   estimateConstantIslandSize(const BinaryFunction *OnBehalfOf = nullptr) const {
     if (!Islands)
@@ -2074,9 +2081,13 @@ public:
       Size += NextMarker - *DataIter;
     }
 
-    if (!OnBehalfOf)
-      for (BinaryFunction *ExternalFunc : Islands->Dependency)
+    if (!OnBehalfOf) {
+      for (BinaryFunction *ExternalFunc : Islands->Dependency) {
+        Size = alignTo(Size, ExternalFunc->getConstantIslandAlignment());
         Size += ExternalFunc->estimateConstantIslandSize(this);
+      }
+    }
+
     return Size;
   }
 
diff --git a/bolt/lib/Core/BinaryEmitter.cpp b/bolt/lib/Core/BinaryEmitter.cpp
index b672a801a360..914549651fe7 100644
--- a/bolt/lib/Core/BinaryEmitter.cpp
+++ b/bolt/lib/Core/BinaryEmitter.cpp
@@ -500,6 +500,13 @@ void BinaryEmitter::emitConstantIslands(BinaryFunction &BF, bool EmitColdPart,
   if (Islands.DataOffsets.empty() && Islands.Dependency.empty())
     return;
 
+  // AArch64 requires CI to be aligned to 8 bytes due to access instructions
+  // restrictions. E.g. the ldr with imm, where imm must be aligned to 8 bytes.
+  const uint16_t Alignment = OnBehalfOf
+                                 ? OnBehalfOf->getConstantIslandAlignment()
+                                 : BF.getConstantIslandAlignment();
+  Streamer.emitCodeAlignment(Alignment, &*BC.STI);
+
   if (!OnBehalfOf) {
     if (!EmitColdPart)
       Streamer.emitLabel(BF.getFunctionConstantIslandLabel());
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 355f35469982..56b50cbdd8cd 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -308,6 +308,7 @@ uint64_t LongJmpPass::tentativeLayoutRelocColdPart(
     LLVM_DEBUG(dbgs() << Func->getPrintName() << " cold tentative: "
                       << Twine::utohexstr(DotAddress) << "\n");
     DotAddress += Func->estimateColdSize();
+    DotAddress = alignTo(DotAddress, Func->getConstantIslandAlignment());
     DotAddress += Func->estimateConstantIslandSize();
   }
   return DotAddress;
@@ -364,6 +365,8 @@ uint64_t LongJmpPass::tentativeLayoutRelocMode(
       DotAddress += Func->estimateSize();
     else
       DotAddress += Func->estimateHotSize();
+
+    DotAddress = alignTo(DotAddress, Func->getConstantIslandAlignment());
     DotAddress += Func->estimateConstantIslandSize();
     ++CurrentIndex;
   }
diff --git a/bolt/test/AArch64/constant-island-alignment.s b/bolt/test/AArch64/constant-island-alignment.s
new file mode 100644
index 000000000000..faf7b0451cd1
--- /dev/null
+++ b/bolt/test/AArch64/constant-island-alignment.s
@@ -0,0 +1,38 @@
+// This test checks that the constant island is aligned after BOLT tool.
+// In case the nop before .Lci will be removed the pointer to exit function
+// won't be alinged and the test will fail.
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: %clang %cflags -fPIC -pie %t.o -o %t.exe -Wl,-q \
+# RUN:    -nostartfiles -nodefaultlibs -lc
+# RUN: llvm-bolt %t.exe -o %t.bolt -use-old-text=0 -lite=0 -trap-old-code
+# RUN: llvm-objdump -d --disassemble-symbols='$d' %t.bolt | FileCheck %s
+
+.text
+.align 4
+.global
+.type dummy, %function
+dummy:
+  add x0, x0, #1
+  ret
+
+.global
+.type exitOk, %function
+exitOk:
+  mov x0, #0
+  bl exit
+
+.global _start
+.type _start, %function
+_start:
+  adrp x0, .Lci
+  ldr x0, [x0, #:lo12:.Lci]
+  blr x0
+  mov x1, #1
+  bl exit
+  nop
+# CHECK: {{0|8}} <$d>:
+.Lci:
+  .xword exitOk
+  .xword 0
diff --git a/bolt/test/runtime/AArch64/adrrelaxationpass.s b/bolt/test/runtime/AArch64/adrrelaxationpass.s
index dc7b49195912..1abb3ee19b73 100644
--- a/bolt/test/runtime/AArch64/adrrelaxationpass.s
+++ b/bolt/test/runtime/AArch64/adrrelaxationpass.s
@@ -41,9 +41,9 @@ br:
   .word 0xff
 
 # CHECK: <main>:
-# CHECK-NEXT: adr x0, #28
+# CHECK-NEXT: adr x0, #{{[0-9][0-9]*}}
 # CHECK-NEXT: adrp x1, 0x{{[1-8a-f][0-9a-f]*}}
 # CHECK-NEXT: add x1, x1, #{{[1-8a-f][0-9a-f]*}}
 # CHECK-NEXT: adrp x2, 0x{{[1-8a-f][0-9a-f]*}}
 # CHECK-NEXT: add x2, x2, #{{[1-8a-f][0-9a-f]*}}
-# CHECK-NEXT: adr x3, #4
+# CHECK-NEXT: adr x3, #{{[0-9][0-9]*}}