[Jump Threading] Unfold a select instruction that feeds a switch statement via a phi node

[amehsan]

Summary

Currently when a select has a constant value in one branch and the select feeds a conditional branch (via a compare/ phi and compare) we unfold the select statement. This results in threading the conditional branch later on. Similar opportunity exists when a select (with a constant in one branch) feeds a switch (via a phi node). The patch unfolds select under this condition. A testcase is provided.

The patch is tested on two out of tree targets with internal benchmarks and also with testsuite on X86. There are several cases where this code unfolds a select in testsuite (when testing on X86) but none resulted in performance improvement or degradation.

[brzycki]

Hi @amehsan , I think the patch looks mostly good but for a few small comments or variable names to enhance readability.

include/llvm/Transforms/Scalar/JumpThreading.h

146

This should probably be written as:

bool TryToUnfoldSelect(SwitchInst *SI, BasicBlock *BB);

lib/Transforms/Scalar/JumpThreading.cpp

1176

Style again:

if (SwitchInst *SI = dyn_cast<SwitchInst>(BB->getTerminator()))
  TryToUnfoldSelect(SI, BB);

2393–2459

Please add documentation here describing the requirements for the code shape checks before entering UnfoldSelectInstr. For example, it's not easy to discern Pred has an unconditional branch to BB without looking at the call sites and the function.

2430

Possibly CondPHI instead of SwVar?

2437

SI is a bit confusing to me. I think of it as the top-level select. This might be better named as PredSI.

2441

Could you add a little more to the comment discussing why we don't want to alter the select for the 2nd and 3rd check here?

test/Transforms/JumpThreading/select.ll

437

Is this check sufficient? There is also a change in land.lhs.true where we no longer contain the %spec.select statement and the terminator instruction changes from unconditional to conditional.

[amehsan]

Thanks @brzycki. I believe I have addressed all your comments. Will upload the modified patch shortly.

[brzycki]

Awesome, thanks for doing this @amehsan .

[amehsan]

Thanks @brzycki for careful review. I need to revive my commit access (not used it for two years now), then I will commit this.

[brzycki]

@amehsan gentle ping. Have you resolved your push access?

[amehsan]

In D55264#1345253, @brzycki wrote:

@amehsan gentle ping. Have you resolved your push access?

I responded to this by email but since I don't yet see the response here: Sorry for the delay. Will try to repeat the test-suite run again and commit this week. Since some time has passed, wanted to be on the safe side.

[amehsan]

Committed in rL350931

include/llvm/Transforms/Scalar/JumpThreading.h

bool ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
                               jumpthreading::ConstantPreference Preference,
                               Instruction *CxtI = nullptr);
 
   bool ProcessBranchOnPHI(PHINode *PN);
   bool ProcessBranchOnXOR(BinaryOperator *BO);
   bool ProcessImpliedCondition(BasicBlock *BB);
 
   bool SimplifyPartiallyRedundantLoad(LoadInst *LI);
+  void UnfoldSelectInstr(BasicBlock *Pred, BasicBlock *BB, SelectInst *SI,
+                         PHINode *SIUse, unsigned Idx);
+
   bool TryToUnfoldSelect(CmpInst *CondCmp, BasicBlock *BB);
+  bool TryToUnfoldSelect(SwitchInst *SI, BasicBlock *BB);
   bool TryToUnfoldSelectInCurrBB(BasicBlock *BB);
 
   bool ProcessGuards(BasicBlock *BB);
   bool ThreadGuard(BasicBlock *BB, IntrinsicInst *Guard, BranchInst *BI);
 
 private:
   BasicBlock *SplitBlockPreds(BasicBlock *BB, ArrayRef<BasicBlock *> Preds,
                               const char *Suffix);

lib/Transforms/Scalar/JumpThreading.cpp

if (CondBr && CondConst) {
 
       // We did not manage to simplify this branch, try to see whether
       // CondCmp depends on a known phi-select pattern.
       if (TryToUnfoldSelect(CondCmp, BB))
         return true;
     }
   }
 
+  if (SwitchInst *SI = dyn_cast<SwitchInst>(BB->getTerminator()))
+    TryToUnfoldSelect(SI, BB);
+
   // Check for some cases that are worth simplifying.  Right now we want to look
   // for loads that are used by a switch or by the condition for the branch.  If
   // we see one, check to see if it's partially redundant.  If so, insert a PHI
   // which can then be used to thread the values.
   Value *SimplifyValue = CondInst;
   if (CmpInst *CondCmp = dyn_cast<CmpInst>(SimplifyValue))
     if (isa<Constant>(CondCmp->getOperand(1)))
       SimplifyValue = CondCmp->getOperand(0);
bool JumpThreadingPass::DuplicateCondBranchOnPHIIntoPred(
 
   // Remove the unconditional branch at the end of the PredBB block.
   OldPredBranch->eraseFromParent();
   DTU->applyUpdates(Updates);
 
   ++NumDupes;
   return true;
 }
 
+// Pred is a predecessor of BB with an unconditional branch to BB. SI is
+// a Select instruction in Pred. BB has other predecessors and SI is used in
+// a PHI node in BB. SI has no other use.
+// A new basic block, NewBB, is created and SI is converted to compare and 
+// conditional branch. SI is erased from parent.
+void JumpThreadingPass::UnfoldSelectInstr(BasicBlock *Pred, BasicBlock *BB,
+                                          SelectInst *SI, PHINode *SIUse,
+                                          unsigned Idx) {
+  // Expand the select.
+  //
+  // Pred --
+  //  |    v
+  //  |  NewBB
+  //  |    |
+  //  |-----
+  //  v
+  // BB
+  BranchInst *PredTerm = dyn_cast<BranchInst>(Pred->getTerminator());
+  BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "select.unfold",
+                                         BB->getParent(), BB);
+  // Move the unconditional branch to NewBB.
+  PredTerm->removeFromParent();
+  NewBB->getInstList().insert(NewBB->end(), PredTerm);
+  // Create a conditional branch and update PHI nodes.
+  BranchInst::Create(NewBB, BB, SI->getCondition(), Pred);
+  SIUse->setIncomingValue(Idx, SI->getFalseValue());
+  SIUse->addIncoming(SI->getTrueValue(), NewBB);
+
+  // The select is now dead.
+  SI->eraseFromParent();
+  DTU->applyUpdates({{DominatorTree::Insert, NewBB, BB},
+                    {DominatorTree::Insert, Pred, NewBB}});
+
+  // Update any other PHI nodes in BB.
+  for (BasicBlock::iterator BI = BB->begin();
+       PHINode *Phi = dyn_cast<PHINode>(BI); ++BI)
+    if (Phi != SIUse)
+      Phi->addIncoming(Phi->getIncomingValueForBlock(Pred), NewBB);
+}
+
+bool JumpThreadingPass::TryToUnfoldSelect(SwitchInst *SI, BasicBlock *BB) {
+  PHINode *CondPHI = dyn_cast<PHINode>(SI->getCondition());
+
+  if (!CondPHI || CondPHI->getParent() != BB)
+    return false;
+
+  for (unsigned I = 0, E = CondPHI->getNumIncomingValues(); I != E; ++I) {
+    BasicBlock *Pred = CondPHI->getIncomingBlock(I);
+    SelectInst *PredSI = dyn_cast<SelectInst>(CondPHI->getIncomingValue(I));
+
+    // The second and third condition can be potentially relaxed. Currently
+    // the conditions help to simplify the code and allow us to reuse existing
+    // code, developed for TryToUnfoldSelect(CmpInst *, BasicBlock *)
+    if (!PredSI || PredSI->getParent() != Pred || !PredSI->hasOneUse())
+      continue;
+
+    BranchInst *PredTerm = dyn_cast<BranchInst>(Pred->getTerminator());
+    if (!PredTerm || !PredTerm->isUnconditional())
+      continue;
+
+    UnfoldSelectInstr(Pred, BB, PredSI, CondPHI, I);
+    return true;
+  }
+  return false;
+}
+
 /// TryToUnfoldSelect - Look for blocks of the form
 /// bb1:
 ///   %a = select
 ///   br bb2
 ///
 /// bb2:
 ///   %p = phi [%a, %bb1] ...
 ///   %c = icmp %p
LazyValueInfo::Tristate LHSFolds =
         LVI->getPredicateOnEdge(CondCmp->getPredicate(), SI->getOperand(1),
                                 CondRHS, Pred, BB, CondCmp);
     LazyValueInfo::Tristate RHSFolds =
         LVI->getPredicateOnEdge(CondCmp->getPredicate(), SI->getOperand(2),
                                 CondRHS, Pred, BB, CondCmp);
     if ((LHSFolds != LazyValueInfo::Unknown ||
          RHSFolds != LazyValueInfo::Unknown) &&
         LHSFolds != RHSFolds) {
-      // Expand the select.
-      //
-      // Pred --
-      //  |    v
-      //  |  NewBB
-      //  |    |
-      //  |-----
-      //  v
-      // BB
-      BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "select.unfold",
-                                             BB->getParent(), BB);
-      // Move the unconditional branch to NewBB.
-      PredTerm->removeFromParent();
-      NewBB->getInstList().insert(NewBB->end(), PredTerm);
-      // Create a conditional branch and update PHI nodes.
-      BranchInst::Create(NewBB, BB, SI->getCondition(), Pred);
-      CondLHS->setIncomingValue(I, SI->getFalseValue());
-      CondLHS->addIncoming(SI->getTrueValue(), NewBB);
-      // The select is now dead.
-      SI->eraseFromParent();
-
-      DTU->applyUpdates({{DominatorTree::Insert, NewBB, BB},
-                         {DominatorTree::Insert, Pred, NewBB}});
-      // Update any other PHI nodes in BB.
-      for (BasicBlock::iterator BI = BB->begin();
-           PHINode *Phi = dyn_cast<PHINode>(BI); ++BI)
-        if (Phi != CondLHS)
-          Phi->addIncoming(Phi->getIncomingValueForBlock(Pred), NewBB);
+      UnfoldSelectInstr(Pred, BB, SI, CondLHS, I);
       return true;
     }
   }
   return false;
 }
 
 /// TryToUnfoldSelectInCurrBB - Look for PHI/Select or PHI/CMP/Select in the
 /// same BB in the form

test/Transforms/JumpThreading/select.ll

 
 ; CHECK-LABEL: @unfold5
 ; CHECK: br i1 %cmp.i, label %.exit, label %cond.false.i
 ; CHECK: br i1 %cmp4.i, label %.exit, label %cond.false.6.i
 ; CHECK: br i1 %cmp8.i, label %.exit, label %cond.false.10.i
 ; CHECK: br i1 %cmp13.i, label %.exit, label %cond.false.15.i
 ; CHECK: br label %.exit
 }
+
+; When a select has a constant operand in one branch, and it feeds a phi node
+; and the phi node feeds a switch we unfold the select
+define void @test_func(i32* nocapture readonly %a, i32* nocapture readonly %b, i32* nocapture readonly %c, i32 %n) local_unnamed_addr #0 {
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %sw.default, %entry
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %sw.default ]
+  %cmp = icmp slt i32 %i.0, %n
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond
+  ret void
+
+for.body:                                         ; preds = %for.cond
+  %0 = zext i32 %i.0 to i64
+  %arrayidx = getelementptr inbounds i32, i32* %a, i64 %0
+  %1 = load i32, i32* %arrayidx, align 4
+  %cmp1 = icmp eq i32 %1, 4
+  br i1 %cmp1, label %land.lhs.true, label %if.end
+
+land.lhs.true:                                    ; preds = %for.body
+  %arrayidx3 = getelementptr inbounds i32, i32* %b, i64 %0
+  %2 = load i32, i32* %arrayidx3, align 4
+  %arrayidx5 = getelementptr inbounds i32, i32* %c, i64 %0
+  %3 = load i32, i32* %arrayidx5, align 4
+  %cmp6 = icmp eq i32 %2, %3
+  %spec.select = select i1 %cmp6, i32 2, i32 4
+  br label %if.end
+
+if.end:                                           ; preds = %land.lhs.true, %for.body
+  %local_var.0 = phi i32 [ %1, %for.body ], [ %spec.select, %land.lhs.true ]
+  switch i32 %local_var.0, label %sw.default [
+    i32 2, label %sw.bb
+    i32 4, label %sw.bb7
+    i32 5, label %sw.bb8
+    i32 7, label %sw.bb9
+  ]
+
+sw.bb:                                            ; preds = %if.end
+  call void @foo()
+  br label %sw.bb7
+
+sw.bb7:                                           ; preds = %if.end, %sw.bb
+  call void @bar()
+  br label %sw.bb8
+
+sw.bb8:                                           ; preds = %if.end, %sw.bb7
+  call void @baz()
+  br label %sw.bb9
+
+sw.bb9:                                           ; preds = %if.end, %sw.bb8
+  call void @quux()
+  br label %sw.default
+
+sw.default:                                       ; preds = %if.end, %sw.bb9
+  call void @baz()
+  %inc = add nuw nsw i32 %i.0, 1
+  br label %for.cond
+
+; CHECK-LABEL: @test_func(
+; CHECK: [[REG:%[0-9]+]] = load
+; CHECK-NOT: select
+; CHECK: br i1
+; CHECK-NOT: select
+; CHECK: br i1 {{.*}}, label [[DEST1:%.*]], label [[DEST2:%.*]]
+
+; The following line checks existence of a phi node, and makes sure
+; it only has one incoming value. To do this, we check every '%'. Note
+; that REG and REG2 each contain one '%;. There is another one in the
+; beginning of the incoming block name. After that there should be no other '%'.
+
+; CHECK: [[REG2:%.*]] = phi i32 {{[^%]*}}[[REG]]{{[^%]*%[^%]*}}
+; CHECK: switch i32 [[REG2]]
+; CHECK: i32 2, label [[DEST1]]
+; CHECK: i32 4, label [[DEST2]]
+}