[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 ↗

(On Diff #102199)

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 ↗

(On Diff #102199)

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 ↗

(On Diff #102199)

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 ↗

(On Diff #102199)

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 ↗

(On Diff #102199)

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

[mssimpso]

Thanks Mathias!

lib/Target/AArch64/AArch64ConditionalCompares.cpp

590 ↗

(On Diff #102199)

Will do.

llvm/trunk/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
+    // successor is Tail. We need to increase the successor probability for
+    // Tail to account for the CmpBB path we removed.
+    //
+    // Pr(Tail|Head) += Pr(CmpBB|Head) * Pr(Tail|CmpBB).
+    assert(*Head->succ_begin() == Tail && "Head successor is not Tail");
+    BranchProbability Head2Tail = MBPI->getEdgeProbability(Head, Tail);
+    Head->setSuccProbability(Head->succ_begin(),
+                             Head2Tail + Head2CmpBB * CmpBB2Tail);
+
+    // We will transfer successors of CmpBB to Head in a moment without
+    // normalizing the successor 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;
 }

llvm/trunk/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

llvm/trunk/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:
+...