[AArch64] Update successor probabilities after ccmp-conversion

[mssimpso]

Summary

This patch modifies the conditional compares pass so that it keeps the successor probabilities of the head block up-to-date after the conversion. Previously, the successor probabilities were being normalized to a uniform distribution, even though they may have been heavily biased prior to the conversion (e.g., if one of the edges was the back edge of a loop). This loss of information affected passes later in the pipeline.

[mssimpso]

Ping.

[MatzeB]

lib/Target/AArch64/AArch64ConditionalCompares.cpp

143

Huh, we can use a const instance of this to change the probabilities... Looking at MachineBranchProbabilityInfo I wonder why this was implemented as a pass and not just a set of utility functions. Anyway just wanted to write this down, it isn't really relevant for the review at hand.

590

I cannot convince myself right now that it is always succ_begin() that points to Tail. Wouldn't we need to search for the right successor?

[mssimpso]

Thanks for the comments Matthias! Responses are inline.

lib/Target/AArch64/AArch64ConditionalCompares.cpp

143

I agree, this was a bit confusing to me. My understanding is that MachineBranchProbabilityInfo is required for reading off the successor probabilities, but that the setting of the probabilities must be done via MachineBasicBlock.

590

Head is required to have exactly two successors, CmpBB and Tail (see the beginning part of canConvert). CmpBB is removed as a successor of Head a few lines up (577), so succ_begin() can only point to Tail here.

[mssimpso]

Ping.

[MatzeB]

LGTM

lib/Target/AArch64/AArch64ConditionalCompares.cpp

590

Ok, maybe check with an assert() to make life easier for readers?

[mssimpso]

Thanks Mathias!

lib/Target/AArch64/AArch64ConditionalCompares.cpp

590

Will do.

lib/Target/AArch64/AArch64ConditionalCompares.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "AArch64.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/MachineTraceMetrics.h"
 #include "llvm/CodeGen/Passes.h"

 // FIXME: Handle PHIs in Tail by turning them into selects (if-conversion).
 
 namespace {
 class SSACCmpConv {
   MachineFunction *MF;
   const TargetInstrInfo *TII;
   const TargetRegisterInfo *TRI;
   MachineRegisterInfo *MRI;
+  const MachineBranchProbabilityInfo *MBPI;
 
 public:
   /// The first block containing a conditional branch, dominating everything
   /// else.
   MachineBasicBlock *Head;
 
   /// The block containing cmp+br.cond with a successor shared with Head.
   MachineBasicBlock *CmpBB;
private:
   MachineInstr *findConvertibleCompare(MachineBasicBlock *MBB);
 
   /// Return true if all non-terminator instructions in MBB can be safely
   /// speculated.
   bool canSpeculateInstrs(MachineBasicBlock *MBB, const MachineInstr *CmpMI);
 
 public:
   /// runOnMachineFunction - Initialize per-function data structures.
-  void runOnMachineFunction(MachineFunction &MF) {
+  void runOnMachineFunction(MachineFunction &MF,
+                            const MachineBranchProbabilityInfo *MBPI) {
     this->MF = &MF;
+    this->MBPI = MBPI;
     TII = MF.getSubtarget().getInstrInfo();
     TRI = MF.getSubtarget().getRegisterInfo();
     MRI = &MF.getRegInfo();
   }
 
   /// If the sub-CFG headed by MBB can be cmp-converted, initialize the
   /// internal state, and return true.
   bool canConvert(MachineBasicBlock *MBB);

 
 void SSACCmpConv::convert(SmallVectorImpl<MachineBasicBlock *> &RemovedBlocks) {
   DEBUG(dbgs() << "Merging BB#" << CmpBB->getNumber() << " into BB#"
                << Head->getNumber() << ":\n" << *CmpBB);
 
   // All CmpBB instructions are moved into Head, and CmpBB is deleted.
   // Update the CFG first.
   updateTailPHIs();
-  Head->removeSuccessor(CmpBB, true);
-  CmpBB->removeSuccessor(Tail, true);
+
+  // Save successor probabilties before removing CmpBB and Tail from their
+  // parents.
+  BranchProbability Head2CmpBB = MBPI->getEdgeProbability(Head, CmpBB);
+  BranchProbability CmpBB2Tail = MBPI->getEdgeProbability(CmpBB, Tail);
+
+  Head->removeSuccessor(CmpBB);
+  CmpBB->removeSuccessor(Tail);
+
+  // If Head and CmpBB had successor probabilties, udpate the probabilities to
+  // reflect the ccmp-conversion.
+  if (Head->hasSuccessorProbabilities() && CmpBB->hasSuccessorProbabilities()) {
+
+    // Head is allowed two successors. We've removed CmpBB, so the remaining
+    // sucessor is Tail. We need to increase the sucessor probability for Tail
+    // to account for the CmpBB path we removed.
+    //
+    // Pr(Tail|Head) += Pr(CmpBB|Head) * Pr(Tail|CmpBB).
+    BranchProbability Head2Tail = MBPI->getEdgeProbability(Head, Tail);
+    Head->setSuccProbability(Head->succ_begin(),
+                             Head2Tail + Head2CmpBB * CmpBB2Tail);
+
+    // We will transfer sucessors of CmpBB to Head in a moment without
+    // normalizing the sucessor probabilities. Set the successor probabilites
+    // before doing so.
+    //
+    // Pr(I|Head) = Pr(CmpBB|Head) * Pr(I|CmpBB).
+    for (auto I = CmpBB->succ_begin(), E = CmpBB->succ_end(); I != E; ++I) {
+      BranchProbability CmpBB2I = MBPI->getEdgeProbability(CmpBB, *I);
+      CmpBB->setSuccProbability(I, Head2CmpBB * CmpBB2I);
+    }
+  }
+
   Head->transferSuccessorsAndUpdatePHIs(CmpBB);
   DebugLoc TermDL = Head->getFirstTerminator()->getDebugLoc();
   TII->removeBranch(*Head);
 
   // If the Head terminator was one of the cbz / tbz branches with built-in
   // compare, we need to insert an explicit compare instruction in its place.
   if (HeadCond[0].getImm() == -1) {
     ++NumCompBranches;

 }
 
 //===----------------------------------------------------------------------===//
 //                       AArch64ConditionalCompares Pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 class AArch64ConditionalCompares : public MachineFunctionPass {
+  const MachineBranchProbabilityInfo *MBPI;
   const TargetInstrInfo *TII;
   const TargetRegisterInfo *TRI;
   MCSchedModel SchedModel;
   // Does the proceeded function has Oz attribute.
   bool MinSize;
   MachineRegisterInfo *MRI;
   MachineDominatorTree *DomTree;
   MachineLoopInfo *Loops;
private:
   bool shouldConvert();
 };
 } // end anonymous namespace
 
 char AArch64ConditionalCompares::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AArch64ConditionalCompares, "aarch64-ccmp",
                       "AArch64 CCMP Pass", false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineBranchProbabilityInfo)
 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
 INITIALIZE_PASS_DEPENDENCY(MachineTraceMetrics)
 INITIALIZE_PASS_END(AArch64ConditionalCompares, "aarch64-ccmp",
                     "AArch64 CCMP Pass", false, false)
 
 FunctionPass *llvm::createAArch64ConditionalCompares() {
   return new AArch64ConditionalCompares();
 }
 
 void AArch64ConditionalCompares::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequired<MachineBranchProbabilityInfo>();
   AU.addRequired<MachineDominatorTree>();
   AU.addPreserved<MachineDominatorTree>();
   AU.addRequired<MachineLoopInfo>();
   AU.addPreserved<MachineLoopInfo>();
   AU.addRequired<MachineTraceMetrics>();
   AU.addPreserved<MachineTraceMetrics>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
if (skipFunction(*MF.getFunction()))
     return false;
 
   TII = MF.getSubtarget().getInstrInfo();
   TRI = MF.getSubtarget().getRegisterInfo();
   SchedModel = MF.getSubtarget().getSchedModel();
   MRI = &MF.getRegInfo();
   DomTree = &getAnalysis<MachineDominatorTree>();
   Loops = getAnalysisIfAvailable<MachineLoopInfo>();
+  MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   Traces = &getAnalysis<MachineTraceMetrics>();
   MinInstr = nullptr;
   MinSize = MF.getFunction()->optForMinSize();
 
   bool Changed = false;
-  CmpConv.runOnMachineFunction(MF);
+  CmpConv.runOnMachineFunction(MF, MBPI);
 
   // Visit blocks in dominator tree pre-order. The pre-order enables multiple
   // cmp-conversions from the same head block.
   // Note that updateDomTree() modifies the children of the DomTree node
   // currently being visited. The df_iterator supports that; it doesn't look at
   // child_begin() / child_end() until after a node has been visited.
   for (auto *I : depth_first(DomTree))
     if (tryConvert(I->getBlock()))
       Changed = true;
 
   return Changed;
 }

test/CodeGen/AArch64/arm64-ccmp.ll

 
 ; Speculatively execute division by zero.
 ; The sdiv/udiv instructions do not trap when the divisor is zero, so they are
 ; safe to speculate.
 ; CHECK-LABEL: speculate_division:
 ; CHECK: cmp w0, #1
 ; CHECK: sdiv [[DIVRES:w[0-9]+]], w1, w0
 ; CHECK: ccmp [[DIVRES]], #16, #0, ge
-; CHECK: b.gt [[BLOCK:LBB[0-9_]+]]
-; CHECK: bl _foo
+; CHECK: b.le [[BLOCK:LBB[0-9_]+]]
 ; CHECK: [[BLOCK]]:
+; CHECK: bl _foo
 ; CHECK: orr w0, wzr, #0x7
 define i32 @speculate_division(i32 %a, i32 %b) nounwind ssp {
 entry:
   %cmp = icmp sgt i32 %a, 0
   br i1 %cmp, label %land.lhs.true, label %if.end
 
 land.lhs.true:
   %div = sdiv i32 %b, %a
   %cmp1 = icmp slt i32 %div, 17
   br i1 %cmp1, label %if.then, label %if.end
 
 if.then:
   %call = tail call i32 @foo() nounwind
   br label %if.end
 
 if.end:
   ret i32 7
 }
 
 ; Floating point compare.
 ; CHECK: single_fcmp
 ; CHECK: cmp
 ; CHECK-NOT: b.
 ; CHECK: fccmp {{.*}}, #8, ge
-; CHECK: b.lt
+; CHECK: b.ge
 define i32 @single_fcmp(i32 %a, float %b) nounwind ssp {
 entry:
   %cmp = icmp sgt i32 %a, 0
   br i1 %cmp, label %land.lhs.true, label %if.end
 
 land.lhs.true:
   %conv = sitofp i32 %a to float
   %div = fdiv float %b, %conv

test/CodeGen/AArch64/ccmp-successor-probs.mir

+# RUN: llc -o - %s -mtriple=aarch64--linux-gnu -mcpu=falkor -run-pass=aarch64-ccmp | FileCheck %s
+---
+# This test checks that successor probabilties are properly updated after a
+# ccmp-conversion.
+#
+# CHECK-LABEL: name: aarch64-ccmp-successor-probs
+# CHECK:       bb.0:
+# CHECK-NEXT:    successors: %bb.2(0x04000000), %bb.3(0x7c000000)
+# CHECK:         CCMPXr %5, %4, 0, 10, implicit-def %nzcv, implicit %nzcv
+#
+name: aarch64-ccmp-successor-probs
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: gpr64 }
+  - { id: 3, class: gpr64 }
+  - { id: 4, class: gpr64 }
+  - { id: 5, class: gpr64 }
+  - { id: 6, class: gpr64 }
+  - { id: 7, class: gpr64 }
+body : |
+  bb.0:
+    successors: %bb.1(0x7e000000), %bb.2(0x02000000)
+
+    %0 = LDRXui killed %x0, 69
+    %1 = COPY %xzr
+    %2 = SUBSXrr %1, %0, implicit-def dead %nzcv
+    %3 = SUBSXri %x1, 1, 0, implicit-def dead %nzcv
+    %4 = COPY %0
+    %5 = COPY %3
+    %6 = SUBSXrr %x1, killed %2, implicit-def %nzcv
+    Bcc 11, %bb.2, implicit %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2(0x02082082), %bb.3(0x7df7df7e)
+
+    %7 = SUBSXrr %5, %4, implicit-def %nzcv
+    Bcc 12, %bb.2, implicit %nzcv
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3(0x80000000)
+
+  bb.3:
+...