[BOLT] Support more than two jump table parents

[aaupov]

Multi-way splitting can cause multiple fragments to access the same jump
table. Relax the assumption that a jump table can only have up to two
parents.

Test Plan: added bolt/test/X86/three-way-split-jt.s

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Multi-way splitting can cause multiple fragments to access the same jump
table. Relax the assumption that a jump table can only have up to two
parents.

Test Plan: added bolt/test/X86/three-way-split-jt.s

---

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

2 Files Affected:

• (modified) bolt/lib/Core/BinaryContext.cpp (+7-5)
• (added) bolt/test/X86/three-way-split-jt.s (+95)

diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 874cdd26ce6ea..88ada1098081d 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -839,9 +839,11 @@ BinaryContext::getOrCreateJumpTable(BinaryFunction &Function, uint64_t Address,
     assert(Address == JT->getAddress() && "unexpected non-empty jump table");
 
     // Prevent associating a jump table to a specific fragment twice.
-    // This simple check arises from the assumption: no more than 2 fragments.
-    if (JT->Parents.size() == 1 && JT->Parents[0] != &Function) {
-      assert(areRelatedFragments(JT->Parents[0], &Function) &&
+    if (!llvm::is_contained(JT->Parents, &Function)) {
+      assert(llvm::all_of(JT->Parents,
+                          [&](const BinaryFunction *BF) {
+                            return areRelatedFragments(&Function, BF);
+                          }) &&
              "cannot re-use jump table of a different function");
       // Duplicate the entry for the parent function for easy access
       JT->Parents.push_back(&Function);
@@ -852,8 +854,8 @@ BinaryContext::getOrCreateJumpTable(BinaryFunction &Function, uint64_t Address,
         JT->print(this->outs());
       }
       Function.JumpTables.emplace(Address, JT);
-      JT->Parents[0]->setHasIndirectTargetToSplitFragment(true);
-      JT->Parents[1]->setHasIndirectTargetToSplitFragment(true);
+      for (BinaryFunction *Parent : JT->Parents)
+        Parent->setHasIndirectTargetToSplitFragment(true);
     }
 
     bool IsJumpTableParent = false;
diff --git a/bolt/test/X86/three-way-split-jt.s b/bolt/test/X86/three-way-split-jt.s
new file mode 100644
index 0000000000000..c4ddc1b9905d7
--- /dev/null
+++ b/bolt/test/X86/three-way-split-jt.s
@@ -0,0 +1,95 @@
+## This reproduces an issue where the function is split into three fragments
+## and all fragments access the same jump table.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.out -v=1 -print-only=main.warm -print-cfg  2>&1 | FileCheck %s
+
+# CHECK-DAG: BOLT-INFO: marking main.warm as a fragment of main
+# CHECK-DAG: BOLT-INFO: marking main.cold as a fragment of main
+# CHECK-DAG: BOLT-INFO: processing main.warm as a sibling of non-ignored function
+# CHECK-DAG: BOLT-INFO: processing main.cold as a sibling of non-ignored function
+# CHECK-DAG: BOLT-WARNING: Ignoring main.cold
+# CHECK-DAG: BOLT-WARNING: Ignoring main.warm
+# CHECK-DAG: BOLT-WARNING: Ignoring main
+# CHECK: BOLT-WARNING: skipped 3 functions due to cold fragments
+
+# CHECK: PIC Jump table JUMP_TABLE for function main, main.warm, main.cold
+# CHECK-NEXT:   0x0000 : __ENTRY_main@0x[[#]]
+# CHECK-NEXT:   0x0004 : __ENTRY_main@0x[[#]]
+# CHECK-NEXT:   0x0008 : __ENTRY_main.cold@0x[[#]]
+# CHECK-NEXT:   0x000c : __ENTRY_main@0x[[#]]
+  .globl main
+  .type main, %function
+  .p2align 2
+main:
+LBB0:
+  andl $0xf, %ecx
+  cmpb $0x4, %cl
+  ## exit through ret
+  ja LBB3
+
+## jump table dispatch, jumping to label indexed by val in %ecx
+LBB1:
+  leaq JUMP_TABLE(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+
+LBB2:
+  xorq %rax, %rax
+LBB3:
+  addq $0x8, %rsp
+  ret
+.size main, .-main
+
+  .globl main.warm
+  .type main.warm, %function
+  .p2align 2
+main.warm:
+LBB20:
+  andl $0xb, %ebx
+  cmpb $0x1, %cl
+  # exit through ret
+  ja LBB23
+
+## jump table dispatch, jumping to label indexed by val in %ecx
+LBB21:
+  leaq JUMP_TABLE(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+
+LBB22:
+  xorq %rax, %rax
+LBB23:
+  addq $0x8, %rsp
+  ret
+.size main.warm, .-main.warm
+
+## cold fragment is only reachable through jump table
+  .globl main.cold
+  .type main.cold, %function
+main.cold:
+  leaq JUMP_TABLE(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+LBB4:
+  callq abort
+.size main.cold, .-main.cold
+
+  .rodata
+## jmp table, entries must be R_X86_64_PC32 relocs
+  .globl JUMP_TABLE
+JUMP_TABLE:
+  .long LBB2-JUMP_TABLE
+  .long LBB3-JUMP_TABLE
+  .long LBB4-JUMP_TABLE
+  .long LBB3-JUMP_TABLE

[rafaelauler]

LG