[SelectOpti][5/5] Optimize select-to-branch transformation

[apostolakis]

Summary

This patch optimizes the transformation of selects to a branch when the heuristics deemed it profitable.
It aggressively sinks eligible instructions to the newly created true/false blocks to prevent their
execution on the common path and interleaves dependence slices to maximize ILP.

Depends on D120232

[apostolakis]

Add tests

[apostolakis]

Rebase and adjust tests

[apostolakis]

Use opt tests instead of llc ones

[apostolakis]

Rebase

[apostolakis]

Fix typo

[apostolakis]

Rebase

[apostolakis]

Rebase

[davidxl]

llvm/lib/CodeGen/SelectOptimize.cpp

313

Is it better to assert that TrueSlices and FalseSlices are not *both* empty? If that is that case, the transformation may be hurtful.

llvm/test/CodeGen/X86/select-optimize.ll

152 ↗

(On Diff #427116)

It is better to add a test case showing sinking for SI group with more than one selects.

[apostolakis]

llvm/lib/CodeGen/SelectOptimize.cpp

313

Sinking is an added optimization on top of the benefit of breaking data dependences.
Even without any sinking, there are many cases where the transformation is still profitable.

For example, consider two dependence chains that include a load.
critical path: .... -> load -> select -> .....
non-critical path \> store
The first one is a long chain and it is the critical path while the second one is a much shorter one.
The load cannot be sunk if the select is converted to a branch since it is used by another instruction (in the other path), but the conversion to a branch breaks a data dependence that reduces the critical path and thus it is still profitable.

llvm/test/CodeGen/X86/select-optimize.ll

152 ↗

(On Diff #427116)

Modified an existing test case with a two-select SI group to include sinking (and an instruction that shouldn't be sunk).

[apostolakis]

Add select-group sinking testing

[davidxl]

llvm/lib/CodeGen/SelectOptimize.cpp

313

Ok. I see that cost analysis already considers the 'sinkability' of the select operands in cost analysis.

620

Should 'ForSinking' be removed? The cost analysis and the actual sinking should be in sync.

[apostolakis]

llvm/lib/CodeGen/SelectOptimize.cpp

620

I think 'ForSinking' should remain. There are instructions that we don't want to sink but could be expensive and add to a dependence chain that might render a conditional move too expensive. For example, a call instruction that might have side effects cannot be sunk, but converting the select to a branch could allow us to avoid blocking on the execution of this call instruction (even if it is on the common path and is scheduled for execution). Thus, it is better to consider the call instruction in the cost analysis.

Additionally, it just happens that one heuristic (base heuristic looking for expensive instructions) shares code with the sinking optimization. Other heuristics (e.g., loop-level, predictable branches) have different cost analyses. If a conversion is deemed profitable we always use the same sinking transformation logic regardless of the used heuristic.

[davidxl]

lgtm

[thakis]

This breaks building on windows: http://45.33.8.238/win/58678/step_4.txt

Please take a look and revert for now if it takes a while to fix.

llvm/lib/CodeGen/SelectOptimize.cpp

 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 #include <algorithm>
 #include <memory>
 #include <queue>

 STATISTIC(NumSelectConvertedExpColdOperand,
           "Number of select groups converted due to expensive cold operand");
 STATISTIC(NumSelectConvertedHighPred,
           "Number of select groups converted due to high-predictability");
 STATISTIC(NumSelectUnPred,
           "Number of select groups not converted due to unpredictability");
 STATISTIC(NumSelectConvertedLoop,
           "Number of select groups converted due to loop-level analysis");
+STATISTIC(NumSelectsConverted, "Number of selects converted");
 
 static cl::opt<unsigned> ColdOperandThreshold(
     "cold-operand-threshold",
     cl::desc("Maximum frequency of path for an operand to be considered cold."),
     cl::init(20), cl::Hidden);
 
 static cl::opt<unsigned>
     LoopGradientGain("loop-gradient-gain",
private:
   void convertProfitableSIGroups(SelectGroups &ProfSIGroups);
   void collectSelectGroups(BasicBlock &BB, SelectGroups &SIGroups);
   void findProfitableSIGroupsBase(SelectGroups &SIGroups,
                                   SelectGroups &ProfSIGroups);
   void findProfitableSIGroupsInnerLoops(const Loop *L, SelectGroups &SIGroups,
                                         SelectGroups &ProfSIGroups);
   bool isConvertToBranchProfitableBase(const SmallVector<SelectInst *, 2> &ASI);
   bool hasExpensiveColdOperand(const SmallVector<SelectInst *, 2> &ASI);
-  void getOneUseSlice(Instruction *I, SmallVector<Instruction *, 2> &Slice);
+  void getOneUseSlice(Instruction *I, std::stack<Instruction *> &Slice,
+                      bool ForSinking = false);
   bool isSelectHighlyPredictable(const SelectInst *SI);
   bool checkLoopHeuristics(CostInfo LoopDepth[2]);
   bool computeLoopCosts(const Loop *L, const SelectGroups &SIGroups,
                         DenseMap<const Instruction *, CostInfo> &InstCostMap,
                         CostInfo LoopCost[2]);
   SmallPtrSet<const Instruction *, 2> getSIset(const SelectGroups &SIGroups);
   Optional<uint64_t> computeInstCost(const Instruction *I);
   uint64_t getMispredictionCost(const SelectInst *SI, const uint64_t CondCost);
for (Loop *L : Loops) {
     SelectGroups SIGroups;
     for (BasicBlock *BB : L->getBlocks())
       collectSelectGroups(*BB, SIGroups);
 
     findProfitableSIGroupsInnerLoops(L, SIGroups, ProfSIGroups);
   }
 }
 
+/// If \p isTrue is true, return the true value of \p SI, otherwise return
+/// false value of \p SI. If the true/false value of \p SI is defined by any
+/// select instructions in \p Selects, look through the defining select
+/// instruction until the true/false value is not defined in \p Selects.
+static Value *
+getTrueOrFalseValue(SelectInst *SI, bool isTrue,
+                    const SmallPtrSet<const Instruction *, 2> &Selects) {
+  Value *V = nullptr;
+  for (SelectInst *DefSI = SI; DefSI != nullptr && Selects.count(DefSI);
+       DefSI = dyn_cast<SelectInst>(V)) {
+    assert(DefSI->getCondition() == SI->getCondition() &&
+           "The condition of DefSI does not match with SI");
+    V = (isTrue ? DefSI->getTrueValue() : DefSI->getFalseValue());
+  }
+  assert(V && "Failed to get select true/false value");
+  return V;
+}
+
 void SelectOptimize::convertProfitableSIGroups(SelectGroups &ProfSIGroups) {
-  llvm_unreachable("Unimplemented");
+  for (SelectGroup &ASI : ProfSIGroups) {
+    // The code transformation here is a modified version of the sinking
+    // transformation in CodeGenPrepare::optimizeSelectInst with a more
+    // aggressive strategy of which instructions to sink.
+    //
+    // TODO: eliminate the redundancy by removing
+    // CodeGenPrepare::optimizeSelectInst and optimizing here selects for all
+    // cases (with and without profile information).
+
+    // Transform a sequence like this:
+    //    start:
+    //       %cmp = cmp uge i32 %a, %b
+    //       %sel = select i1 %cmp, i32 %c, i32 %d
+    //
+    // Into:
+    //    start:
+    //       %cmp = cmp uge i32 %a, %b
+    //       %cmp.frozen = freeze %cmp
+    //       br i1 %cmp.frozen, label %select.true, label %select.false
+    //    select.true:
+    //       br label %select.end
+    //    select.false:
+    //       br label %select.end
+    //    select.end:
+    //       %sel = phi i32 [ %c, %select.true ], [ %d, %select.false ]
+    //
+    // %cmp should be frozen, otherwise it may introduce undefined behavior.
+    // In addition, we may sink instructions that produce %c or %d into the
+    // destination(s) of the new branch.
+    // If the true or false blocks do not contain a sunken instruction, that
+    // block and its branch may be optimized away. In that case, one side of the
+    // first branch will point directly to select.end, and the corresponding PHI
+    // predecessor block will be the start block.
+
+    // Find all the instructions that can be soundly sunk to the true/false
+    // blocks. These are instructions that are computed solely for producing the
+    // operands of the select instructions in the group and can be sunk without
+    // breaking the semantics of the LLVM IR (e.g., cannot sink instructions
+    // with side effects).
+    SmallVector<std::stack<Instruction *>, 2> TrueSlices, FalseSlices;
+    unsigned long maxTrueSliceLen = 0, maxFalseSliceLen = 0;
+    for (SelectInst *SI : ASI) {
+      // For each select, compute the sinkable dependence chains of the true and
+      // false operands.
+      if (auto *TI = dyn_cast<Instruction>(SI->getTrueValue())) {
+        std::stack<Instruction *> TrueSlice;
+        getOneUseSlice(TI, TrueSlice, true);
+        maxTrueSliceLen = std::max(maxTrueSliceLen, TrueSlice.size());
+        TrueSlices.push_back(TrueSlice);
+      }
+      if (auto *FI = dyn_cast<Instruction>(SI->getFalseValue())) {
+        std::stack<Instruction *> FalseSlice;
+        getOneUseSlice(FI, FalseSlice, true);
+        maxFalseSliceLen = std::max(maxFalseSliceLen, FalseSlice.size());
+        FalseSlices.push_back(FalseSlice);
+      }
+    }
+    // In the case of multiple select instructions in the same group, the order
+    // of non-dependent instructions (instructions of different dependence
+    // slices) in the true/false blocks appears to affect performance.
+    // Interleaving the slices seems to experimentally be the optimal approach.
+    // This interleaving scheduling allows for more ILP (with a natural downside
+    // of increasing a bit register pressure) compared to a simple ordering of
+    // one whole chain after another. One would expect that this ordering would
+    // not matter since the scheduling in the backend of the compiler  would
+    // take care of it, but apparently the scheduler fails to deliver optimal
+    // ILP with a naive ordering here.
+    SmallVector<Instruction *, 2> TrueSlicesInterleaved, FalseSlicesInterleaved;
+    for (unsigned long IS = 0; IS < maxTrueSliceLen; ++IS) {
+      for (auto &S : TrueSlices) {
+        if (!S.empty()) {
+          TrueSlicesInterleaved.push_back(S.top());
+          S.pop();
+        }
+      }
+    }
+    for (unsigned long IS = 0; IS < maxFalseSliceLen; ++IS) {
+      for (auto &S : FalseSlices) {
+        if (!S.empty()) {
+          FalseSlicesInterleaved.push_back(S.top());
+          S.pop();
+        }
+      }
+    }
+
+    // We split the block containing the select(s) into two blocks.
+    SelectInst *SI = ASI.front();
+    SelectInst *LastSI = ASI.back();
+    BasicBlock *StartBlock = SI->getParent();
+    BasicBlock::iterator SplitPt = ++(BasicBlock::iterator(LastSI));
+    BasicBlock *EndBlock = StartBlock->splitBasicBlock(SplitPt, "select.end");
+    BFI->setBlockFreq(EndBlock, BFI->getBlockFreq(StartBlock).getFrequency());
+    // Delete the unconditional branch that was just created by the split.
+    StartBlock->getTerminator()->eraseFromParent();
+
+    // These are the new basic blocks for the conditional branch.
+    // At least one will become an actual new basic block.
+    BasicBlock *TrueBlock = nullptr, *FalseBlock = nullptr;
+    BranchInst *TrueBranch = nullptr, *FalseBranch = nullptr;
+    if (!TrueSlicesInterleaved.empty()) {
+      TrueBlock = BasicBlock::Create(LastSI->getContext(), "select.true.sink",
+                                     EndBlock->getParent(), EndBlock);
+      TrueBranch = BranchInst::Create(EndBlock, TrueBlock);
+      TrueBranch->setDebugLoc(LastSI->getDebugLoc());
+      for (Instruction *TrueInst : TrueSlicesInterleaved)
+        TrueInst->moveBefore(TrueBranch);
+    }
+    if (!FalseSlicesInterleaved.empty()) {
+      FalseBlock = BasicBlock::Create(LastSI->getContext(), "select.false.sink",
+                                      EndBlock->getParent(), EndBlock);
+      FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
+      FalseBranch->setDebugLoc(LastSI->getDebugLoc());
+      for (Instruction *FalseInst : FalseSlicesInterleaved)
+        FalseInst->moveBefore(FalseBranch);
+    }
+    // If there was nothing to sink, then arbitrarily choose the 'false' side
+    // for a new input value to the PHI.
+    if (TrueBlock == FalseBlock) {
+      assert(TrueBlock == nullptr &&
+             "Unexpected basic block transform while optimizing select");
+
+      FalseBlock = BasicBlock::Create(SI->getContext(), "select.false",
+                                      EndBlock->getParent(), EndBlock);
+      auto *FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
+      FalseBranch->setDebugLoc(SI->getDebugLoc());
+    }
+
+    // Insert the real conditional branch based on the original condition.
+    // If we did not create a new block for one of the 'true' or 'false' paths
+    // of the condition, it means that side of the branch goes to the end block
+    // directly and the path originates from the start block from the point of
+    // view of the new PHI.
+    BasicBlock *TT, *FT;
+    if (TrueBlock == nullptr) {
+      TT = EndBlock;
+      FT = FalseBlock;
+      TrueBlock = StartBlock;
+    } else if (FalseBlock == nullptr) {
+      TT = TrueBlock;
+      FT = EndBlock;
+      FalseBlock = StartBlock;
+    } else {
+      TT = TrueBlock;
+      FT = FalseBlock;
+    }
+    IRBuilder<> IB(SI);
+    auto *CondFr =
+        IB.CreateFreeze(SI->getCondition(), SI->getName() + ".frozen");
+    IB.CreateCondBr(CondFr, TT, FT, SI);
+
+    SmallPtrSet<const Instruction *, 2> INS;
+    INS.insert(ASI.begin(), ASI.end());
+    // Use reverse iterator because later select may use the value of the
+    // earlier select, and we need to propagate value through earlier select
+    // to get the PHI operand.
+    for (auto It = ASI.rbegin(); It != ASI.rend(); ++It) {
+      SelectInst *SI = *It;
+      // The select itself is replaced with a PHI Node.
+      PHINode *PN = PHINode::Create(SI->getType(), 2, "", &EndBlock->front());
+      PN->takeName(SI);
+      PN->addIncoming(getTrueOrFalseValue(SI, true, INS), TrueBlock);
+      PN->addIncoming(getTrueOrFalseValue(SI, false, INS), FalseBlock);
+      PN->setDebugLoc(SI->getDebugLoc());
+
+      SI->replaceAllUsesWith(PN);
+      SI->eraseFromParent();
+      INS.erase(SI);
+      ++NumSelectsConverted;
+    }
+  }
 }
 
 void SelectOptimize::collectSelectGroups(BasicBlock &BB,
                                          SelectGroups &SIGroups) {
   BasicBlock::iterator BBIt = BB.begin();
   while (BBIt != BB.end()) {
     Instruction *I = &*BBIt++;
     if (SelectInst *SI = dyn_cast<SelectInst>(I)) {
bool SelectOptimize::hasExpensiveColdOperand(
   // any of the selects of the group.
   for (SelectInst *SI : ASI) {
     Instruction *ColdI = nullptr;
     if (TrueWeight < FalseWeight)
       ColdI = dyn_cast<Instruction>(SI->getTrueValue());
     else
       ColdI = dyn_cast<Instruction>(SI->getFalseValue());
     if (ColdI) {
-      SmallVector<Instruction *, 2> ColdSlice;
+      std::stack<Instruction *> ColdSlice;
       getOneUseSlice(ColdI, ColdSlice);
-      for (auto *ColdII : ColdSlice)
-        if (isExpensiveInst(ColdII))
+      while (!ColdSlice.empty()) {
+        if (isExpensiveInst(ColdSlice.top()))
           return true;
+        ColdSlice.pop();
+      }
     }
   }
   return false;
 }
 
 void SelectOptimize::getOneUseSlice(Instruction *I,
-                                    SmallVector<Instruction *, 2> &Slice) {
+                                    std::stack<Instruction *> &Slice,
+                                    bool ForSinking) {
   std::queue<Instruction *> Worklist;
   Worklist.push(I);
   SmallPtrSet<Instruction *, 2> Visited;
   const Loop *L = LI->getLoopFor(I->getParent());
   while (!Worklist.empty()) {
     Instruction *II = Worklist.front();
     Worklist.pop();
 
     if (Visited.count(II))
       continue;
     Visited.insert(II);
 
     if (!II->hasOneUse())
       continue;
 
-    // Avoid considering instructions in outer or independent loops, or non-loop
-    // instructions (instructions with potentially much less frequency than the
-    // operand of the select).
+    // Cannot soundly sink instructions with side-effects.
+    // Terminator instructions cannot be sunk.
+    // Avoid sinking other select instructions (should be handled separetely).
+    if (ForSinking &&
+        (II->isTerminator() || II->mayHaveSideEffects() || isa<SelectInst>(II)))
+      continue;
+
+    // Avoid considering/sinking instructions from outer or independent loops,
+    // or non-loop instructions (instructions with potentially much less
+    // frequency than the operand of the select).
     const Loop *LII = LI->getLoopFor(II->getParent());
     if (L && L != LII && (!LII || !L->contains(LII)))
       continue;
 
     // Avoid phi nodes that are not part of a loop-carried dependence.
     // Could source from a cold basic block.
     if (isa<PHINode>(II)) {
+      if (ForSinking)
+        continue;
       if (!LII)
         continue;
       if (LII->getHeader() != II->getParent())
         continue;
     }
 
-    Slice.push_back(II);
+    Slice.push(II);
     for (unsigned k = 0; k < II->getNumOperands(); ++k)
       if (auto *OpI = dyn_cast<Instruction>(II->getOperand(k)))
         Worklist.push(OpI);
   }
 }
 
 bool SelectOptimize::isSelectHighlyPredictable(const SelectInst *SI) {
   uint64_t TrueWeight, FalseWeight;