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

[yozhu]

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

[llvmbot]

@llvm/pr-subscribers-bolt

Author: YongKang Zhu (yozhu)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/159486.diff

3 Files Affected:

• (modified) bolt/include/bolt/Core/BinaryFunction.h (+6-6)
• (modified) bolt/lib/Core/BinaryFunction.cpp (+27)
• (modified) bolt/test/AArch64/constant-island-alignment.s (+76-6)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 51b139a15e1a0..7e0e3bff83259 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/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);
   }
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 578a87dc6c09d..46baa120ba53b 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/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.
+  uint64_t MaxAlignment =
+      std::min(uint64_t(1) << llvm::countr_zero(getAddress()),
+               OriginSection->getAlignment());
+  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;
 
diff --git a/bolt/test/AArch64/constant-island-alignment.s b/bolt/test/AArch64/constant-island-alignment.s
index 3ce0df9d4f290..957c4705f5eec 100644
--- a/bolt/test/AArch64/constant-island-alignment.s
+++ b/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

[maksfb]

nit: constify variables that don't change.