[BOLT] Add heuristics to determine constant island's alignment

bolt/include/bolt/Core/BinaryFunction.h

@@ -192,9 +192,6 @@ class BinaryFunction {
 
     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 =
@@ -2114,9 +2111,7 @@ class BinaryFunction {
     return *std::prev(CodeIter) <= *DataIter;
   }
 
-  uint16_t getConstantIslandAlignment() const {
-    return Islands ? Islands->getAlignment() : 1;
-  }
+  uint16_t getConstantIslandAlignment() const;
 
   /// If there is a constant island in the range [StartOffset, EndOffset),
   /// return its address.
@@ -2168,6 +2163,11 @@ class BinaryFunction {
     return Islands && !Islands->DataOffsets.empty();
   }
 
+  /// Return true if the whole function is a constant island.
+  bool isDataObject() const {
+    return Islands && Islands->CodeOffsets.size() == 0;
+  }
+
   bool isStartOfConstantIsland(uint64_t Offset) const {
     return hasConstantIsland() && Islands->DataOffsets.count(Offset);
   }

bolt/lib/Core/BinaryFunction.cpp

@@ -284,6 +284,33 @@ BinaryFunction::getBasicBlockContainingOffset(uint64_t Offset) {
   return (Offset < BB->getOffset() + BB->getOriginalSize()) ? BB : nullptr;
 }
 
+uint16_t BinaryFunction::getConstantIslandAlignment() const {
+  if (Islands == nullptr)
+    return 1;
+
+  // For constant island inside a function, the default 8-byte alignment is
+  // probably good enough.
+  const uint16_t DefaultAlignment = sizeof(uint64_t);
+  if (!isDataObject())
+    return DefaultAlignment;
+
+  // If the constant island itself is a binary function, get its alignment
+  // based on its size, original address, and its owning section's alignment.
+  const uint64_t MaxAlignment =
+      std::min(uint64_t(1) << llvm::countr_zero(getAddress()),
+               OriginSection->getAlignment());
+  const uint64_t MinAlignment =
+      std::max((uint64_t)DefaultAlignment,
+               uint64_t(1) << (63 - llvm::countl_zero(getSize())));
+  uint64_t Alignment = std::min(MinAlignment, MaxAlignment);
+  if (Alignment >> 16) {
+    BC.errs() << "BOLT-ERROR: the constant island's alignment is too big: 0x"
+              << Twine::utohexstr(Alignment) << "\n";
+    exit(1);
+  }
+  return (uint16_t)Alignment;
+}
+
 void BinaryFunction::markUnreachableBlocks() {
   std::stack<BinaryBasicBlock *> Stack;
 

bolt/test/AArch64/constant-island-alignment.s

@@ -1,14 +1,36 @@
 // 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: split-file %s %t
+
+// For the first test case, 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:   %t/xword_align.s -o %t_xa.o
+# RUN: %clang %cflags -fPIC -pie %t_xa.o -o %t_xa.exe -Wl,-q \
 # RUN:    -nostartfiles -nodefaultlibs -Wl,-z,notext
-# 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
+# RUN: llvm-bolt %t_xa.exe -o %t_xa.bolt --use-old-text=0 --lite=0 \
+# RUN:    --trap-old-code
+# RUN: llvm-objdump -d --disassemble-symbols='$d' %t_xa.bolt | FileCheck %s
+
+// For the second and third test cases, we want to set an alignment based
+// on various heuristics.
+
+# RUN: %clang %cflags -pie %t/page_align.s -o %t_pa.exe -Wl,-q \
+# RUN:    -Wl,--init=_foo -Wl,--fini=_foo
+# RUN: llvm-bolt %t_pa.exe -o %t_pa.bolt
+# RUN: llvm-objdump -t %t_pa.exe | grep _const_island
+# RUN: llvm-objdump -t %t_pa.bolt | grep _const_island | FileCheck %s \
+# RUN:    --check-prefix=PAGE
+
+# RUN: %clang %cflags -pie %t/64B_align.s -o %t_64B.exe -Wl,-q \
+# RUN:    -Wl,--init=_foo -Wl,--fini=_foo
+# RUN: llvm-bolt %t_64B.exe -o %t_64B.bolt
+# RUN: llvm-objdump -t %t_64B.exe | grep _const_island
+# RUN: llvm-objdump -t %t_64B.bolt | grep _const_island | FileCheck %s \
+# RUN:    --check-prefix=64BYTE
 
+;--- xword_align.s
 .text
 .align 4
 .global
@@ -36,3 +58,51 @@ _start:
 .Lci:
   .xword exitOk
   .xword 0
+
+;--- page_align.s
+  .text
+  .global _foo
+  .type _foo, %function
+_foo:
+  ret
+
+  .text
+  .global _const_island
+  .align 12
+# PAGE: {{[0-9a-f]*}}000 g
+_const_island:
+  .rept 0x25100
+    .byte 0xbb
+  .endr
+
+  .global _start
+  .type _start, %function
+_start:
+  ret
+
+  # Dummy relocation to force relocation mode
+  .reloc 0, R_AARCH64_NONE
+
+;--- 64B_align.s
+  .text
+  .global _foo
+  .type _foo, %function
+_foo:
+  ret
+
+  .text
+  .global _const_island
+  .align 6
+# 64BYTE: {{[0-9a-f]*}}{{0|4|8|c}}0 g
+_const_island:
+  .rept 0x2048
+    .byte 0xbb
+  .endr
+
+  .global _start
+  .type _start, %function
+_start:
+  ret
+
+  # Dummy relocation to force relocation mode
+  .reloc 0, R_AARCH64_NONE