[BOLT] Consider Code Fragments during regreassign

[hzq]

Summary

During register swapping, the code fragments associated with the
function need to be swapped together (which may be generated during
PGO optimization).

Fix https://github.com/llvm/llvm-project/issues/59730

[rafauler]

Thanks for fixing this, @hzq !

Some suggestions below.

Ideally we would have a test case loosely similar to this one (https://reviews.llvm.org/D138245) that shows that the fix is working as intended. Let me know if you need help writing this.

Finally, be aware that this pass is also experimental because it does not update debug info. So it really depends on how serious you are about keeping debug info accurate.

bolt/include/bolt/Core/BinaryFunction.h

377

1798–1800

bolt/lib/Passes/RegReAssign.cpp

326–328

378–380

[hzq]

Ideally we would have a test case loosely similar to this one (https://reviews.llvm.org/D138245) that shows that the fix is working as intended. Let me know if you need help writing this.

Finally, be aware that this pass is also experimental because it does not update debug info. So it really depends on how serious you are about keeping debug info accurate.

Thanks. will modify and try to write test cases.

[hzq]

Hi, @rafauler , I'm sorry for the delayed response. Could you please help me take a look at this test case to see if it is valid?

[hzq]

update test case and code.

[hzq]

Update.

[rafauler]

Thanks! Mostly formatting nits:

bolt/include/bolt/Core/BinaryFunction.h

377

Since you added a new getter for "getParentFraments", move this "using" to line 369 and change ParentFragments to use FrafmentsSetTy, just like Fragments.

bolt/lib/Passes/RegReAssign.cpp

195–197

nit: drop braces here:

198–202

229

clang-format

230–232

342

clang-format

343–345

bolt/test/X86/reg-reassign-swap-cold.s

12 ↗

(On Diff #546864)

Because this tests runs the produced output, we need to move this test to the folder test/runtime/X86

[hzq]

update+format

[rafauler]

Thanks! LGTM with some modifications I made to the test file available here: https://pastebin.com/3abP1S5G

I modified the test to fail if the optimization is applied just to main, but not to main.cold, and vice-versa. Please commit this version. I also used basic block labels when generating FDATA instead of hardcoded offsets, so the test is resilient to changes in it. I further modified main.cold to be a global symbol, so the symbol is not stripped when we run llvm-strip. We do need to run llvm-strip because I'm not declaring bb labels as temp symbols anymore (.L1, etc), in order to use their names when generating fdata. That's the way we usually write the tests (function names are globals, bb names are locals but non-temp so we don't need to hardcode fdata).

[hzq]

update testcase.

[hzq]

In D141931#4570722, @rafauler wrote:

Thanks! LGTM with some modifications I made to the test file available here: https://pastebin.com/3abP1S5G

I modified the test to fail if the optimization is applied just to main, but not to main.cold, and vice-versa. Please commit this version. I also used basic block labels when generating FDATA instead of hardcoded offsets, so the test is resilient to changes in it. I further modified main.cold to be a global symbol, so the symbol is not stripped when we run llvm-strip. We do need to run llvm-strip because I'm not declaring bb labels as temp symbols anymore (.L1, etc), in order to use their names when generating fdata. That's the way we usually write the tests (function names are globals, bb names are locals but non-temp so we don't need to hardcode fdata).

Thanks very much for your assistance. It has been a great help to me.

bolt/include/bolt/Core/BinaryFunction.h

private:
 
   /// Name for the section this function code should reside in.
   std::string CodeSectionName;
 
   /// Name for the corresponding cold code section.
   std::string ColdCodeSectionName;
 
   /// Parent function fragment for split function fragments.
-  SmallPtrSet<BinaryFunction *, 1> ParentFragments;
+  using FragmentsSetTy = SmallPtrSet<BinaryFunction *, 1>;
+  FragmentsSetTy ParentFragments;
 
   /// Indicate if the function body was folded into another function.
   /// Used by ICF optimization.
   BinaryFunction *FoldedIntoFunction{nullptr};
 
   /// All fragments for a parent function.
-  SmallPtrSet<BinaryFunction *, 1> Fragments;
+  FragmentsSetTy Fragments;
 
   /// The profile data for the number of times the function was executed.
   uint64_t ExecutionCount{COUNT_NO_PROFILE};
 
   /// Profile match ratio.
   float ProfileMatchRatio{0.0f};
 
   /// Raw branch count for this function in the profile.
bool isChildOf(const BinaryFunction &Other) const {
     return ParentFragments.contains(&Other);
   }
 
   /// Returns if this function is a parent of \p Other function.
   bool isParentOf(const BinaryFunction &Other) const {
     return llvm::is_contained(Fragments, &Other);
   }
 
+  /// Return the child fragment form parent function
+  iterator_range<FragmentsSetTy::const_iterator> getFragments() const {
+    return iterator_range<FragmentsSetTy::const_iterator>(Fragments.begin(),
+                                                          Fragments.end());
+  }
+
+  /// Return the parent function for split function fragments.
+  FragmentsSetTy *getParentFragments() { return &ParentFragments; }
+
   /// Returns if this function is a parent or child of \p Other function.
   bool isParentOrChildOf(const BinaryFunction &Other) const {
     return isChildOf(Other) || isParentOf(Other);
   }
 
   /// Set the profile data for the number of times the function was called.
   BinaryFunction &setExecutionCount(uint64_t Count) {
     ExecutionCount = Count;
     return *this;
   }
 
   /// Adjust execution count for the function by a given \p Count. The value
   /// \p Count will be subtracted from the current function count.
   ///
   /// The function will proportionally adjust execution count for all
   /// basic blocks and edges in the control flow graph.
   void adjustExecutionCount(uint64_t Count);

bolt/lib/Passes/RegReAssign.cpp

void RegReAssign::swap(BinaryFunction &Function, MCPhysReg A, MCPhysReg B) {
   }
 }
 
 void RegReAssign::rankRegisters(BinaryFunction &Function) {
   BinaryContext &BC = Function.getBinaryContext();
   std::fill(RegScore.begin(), RegScore.end(), 0);
   std::fill(RankedRegs.begin(), RankedRegs.end(), 0);
 
-  for (BinaryBasicBlock &BB : Function) {
+  auto countRegScore = [&](BinaryBasicBlock &BB) {
     for (MCInst &Inst : BB) {
       const bool CannotUseREX = BC.MIB->cannotUseREX(Inst);
       const MCInstrDesc &Desc = BC.MII->get(Inst.getOpcode());
 
       // Disallow substituitions involving regs in implicit uses lists
       for (MCPhysReg ImplicitUse : Desc.implicit_uses()) {
         const size_t RegEC =
             BC.MIB->getAliases(ImplicitUse, false).find_first();
for (MCInst &Inst : BB) {
           RegScore[RegEC] =
               std::numeric_limits<decltype(RegScore)::value_type>::min();
           continue;
         }
 
         RegScore[RegEC] += BB.getKnownExecutionCount();
       }
     }
+  };
+  for (BinaryBasicBlock &BB : Function)
+    countRegScore(BB);
+
+  for (BinaryFunction *ChildFrag : Function.getFragments()) {
+    for (BinaryBasicBlock &BB : *ChildFrag)
+      countRegScore(BB);
   }
+
   std::iota(RankedRegs.begin(), RankedRegs.end(), 0); // 0, 1, 2, 3...
   llvm::sort(RankedRegs,
              [&](size_t A, size_t B) { return RegScore[A] > RegScore[B]; });
 
   LLVM_DEBUG({
     for (size_t Reg : RankedRegs) {
       if (RegScore[Reg] == 0)
         continue;
       dbgs() << Reg << " ";
       if (RegScore[Reg] > 0)
         dbgs() << BC.MRI->getName(Reg) << ": " << RegScore[Reg] << "\n";
       else
         dbgs() << BC.MRI->getName(Reg) << ": (blacklisted)\n";
     }
   });
 }
 
 void RegReAssign::aggressivePassOverFunction(BinaryFunction &Function) {
   BinaryContext &BC = Function.getBinaryContext();
   rankRegisters(Function);
 
+  // If there is a situation where function:
+  //   A() -> A.cold()
+  //   A.localalias() -> A.cold()
+  // simply swapping these two calls can cause issues.
+  for (BinaryFunction *ChildFrag : Function.getFragments()) {
+    if (ChildFrag->getParentFragments()->size() > 1)
+      return;
+    if (ChildFrag->empty())
+      return;
+  }
+
   // Bail early if our registers are all black listed, before running expensive
   // analysis passes
   bool Bail = true;
   int64_t LowScoreClassic = std::numeric_limits<int64_t>::max();
   for (int J : ClassicRegs.set_bits()) {
     if (RegScore[J] <= 0)
       continue;
     Bail = false;
if (AnyAliasAlive.any()) {
       continue;
     }
 
     // Opportunity detected. Swap.
     LLVM_DEBUG(dbgs() << "\n ** Swapping " << BC.MRI->getName(ClassicReg)
                       << " with " << BC.MRI->getName(ExtReg) << "\n\n");
     swap(Function, ClassicReg, ExtReg);
     FuncsChanged.insert(&Function);
+    for (BinaryFunction *ChildFrag : Function.getFragments()) {
+      swap(*ChildFrag, ClassicReg, ExtReg);
+      FuncsChanged.insert(ChildFrag);
+    }
     ++Begin;
     if (Begin == End)
       break;
     --End;
   }
 }
 
 bool RegReAssign::conservativePassOverFunction(BinaryFunction &Function) {
   BinaryContext &BC = Function.getBinaryContext();
   rankRegisters(Function);
 
+  for (BinaryFunction *ChildFrag : Function.getFragments()) {
+    if (ChildFrag->getParentFragments()->size() > 1)
+      return false;
+    if (ChildFrag->empty())
+      return false;
+  }
+
   // Try swapping R12, R13, R14 or R15 with RBX (we work with all callee-saved
   // regs except RBP)
   MCPhysReg Candidate = 0;
   for (int J : ExtendedCSR.set_bits())
     if (RegScore[J] > RegScore[Candidate])
       Candidate = J;
 
   if (!Candidate || RegScore[Candidate] < 0)
bool RegReAssign::conservativePassOverFunction(BinaryFunction &Function) {
   if (!RBX)
     return false;
 
   LLVM_DEBUG(dbgs() << "\n ** Swapping " << BC.MRI->getName(RBX) << " with "
                     << BC.MRI->getName(Candidate) << "\n\n");
   (void)BC;
   swap(Function, RBX, Candidate);
   FuncsChanged.insert(&Function);
+  for (BinaryFunction *ChildFrag : Function.getFragments()) {
+    swap(*ChildFrag, RBX, Candidate);
+    FuncsChanged.insert(ChildFrag);
+  }
   return true;
 }
 
 void RegReAssign::setupAggressivePass(BinaryContext &BC,
                                       std::map<uint64_t, BinaryFunction> &BFs) {
   setupConservativePass(BC, BFs);
   CG.reset(new BinaryFunctionCallGraph(buildCallGraph(BC)));
   RA.reset(new RegAnalysis(BC, &BFs, &*CG));
void RegReAssign::runOnFunctions(BinaryContext &BC) {
   if (opts::AggressiveReAssign)
     setupAggressivePass(BC, BC.getBinaryFunctions());
   else
     setupConservativePass(BC, BC.getBinaryFunctions());
 
   for (auto &I : BC.getBinaryFunctions()) {
     BinaryFunction &Function = I.second;
 
-    if (!Function.isSimple() || Function.isIgnored())
+    if (!Function.isSimple() || Function.isIgnored() || Function.isFragment())
       continue;
 
     LLVM_DEBUG(dbgs() << "====================================\n");
     LLVM_DEBUG(dbgs() << " - " << Function.getPrintName() << "\n");
     if (!conservativePassOverFunction(Function) && opts::AggressiveReAssign) {
       aggressivePassOverFunction(Function);
       LLVM_DEBUG({
         if (FuncsChanged.count(&Function))

bolt/test/runtime/X86/reg-reassign-swap-cold.s

+# This test case reproduces a bug where, during register swapping,
+# the code fragments associated with the function need to be swapped
+# together (which may be generated during PGO optimization). If not
+# handled properly, optimized binary execution can result in a segmentation fault.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang -no-pie %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.out -data=%t.fdata --reg-reassign | FileCheck %s
+# RUN: %t.out
+
+# CHECK: BOLT-INFO: Reg Reassignment Pass Stats
+# CHECK-NEXT: 2 functions affected.
+  .text
+  .globl  main
+  .globl  main.cold
+  .p2align  4, 0x90
+  .type   main,@function
+  .type   main.cold,@function
+main.cold:
+bb1:
+  cmp     $0x3, %r12
+  jne     bb8
+bb2:
+  jmp     bb4
+main:                                   # @main
+  .cfi_startproc
+# %bb.0:                                # %entry
+  pushq   %rax
+  pushq   %r12
+  pushq   %rbx
+  .cfi_def_cfa_offset 16
+  mov     $0x1,  %r12
+  mov     $0x2,  %rbx
+  add     $0x1,  %r12
+  shr     $0x14, %r12
+  mov     $0x3,  %r12
+bb3:
+  jmp     bb1
+bb4:
+  cmp     $0x3,  %r12
+bb5:
+  jne     bb8
+bb6:
+  xorl    %eax, %eax
+bb7:
+  popq    %rcx
+  popq    %rbx
+  popq    %r12
+  .cfi_def_cfa_offset 8
+  retq
+bb8:
+  mov  $0x1, %rax
+  jmp  bb7
+# FDATA: 1 main.cold #bb2# 1 main #bb4# 0 100
+# FDATA: 1 main #bb5# 1 main #bb6# 0 100
+# FDATA: 1 main #bb3# 1 main.cold 0 0 100
+
+.Lfunc_end0:
+  .size  main, .Lfunc_end0-main
+  .cfi_endproc