[BOLT] Add heuristics to determine constant island's alignment (#159486)

Constant island embedded in text section doesn't have its alignment
information from input binary and we currently set its alignment as
8 bytes. Constant island might be given a much larger alignment due
to performance or other reasons, so this change adds some heuristics
to determine its alignment based on its size, original address from
input binary and its owning section's alignment.

Author: yozhu

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