[DFAJumpThreading] Relax analysis to handle unpredictable initial values

Responding to a feature request from the Rust community:

rust-lang/rust#80630

    void foo(X) {
      for (...)
	switch (X)
	  case A
	    X = B
	  case B
	    X = C
    }

Even though the initial switch value is non-constant, the switch
statement can still be threaded: the initial value will hit the switch
statement but the rest of the state changes will proceed by jumping
unconditionally.

The early predictability check is relaxed to allow unpredictable values
anywhere, but later, after the paths through the switch statement have
been enumerated, no non-constant state values are allowed along the
paths. Any state value not along a path will be an initial switch value,
which can be safely ignored.

Differential Revision: https://reviews.llvm.org/D124394

Author: Alex Zhikhartsev

llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp

@@ -103,6 +103,11 @@ static cl::opt<unsigned> MaxPathLength(
     cl::desc("Max number of blocks searched to find a threading path"),
     cl::Hidden, cl::init(20));
 
+static cl::opt<unsigned> MaxNumPaths(
+    "dfa-max-num-paths",
+    cl::desc("Max number of paths enumerated around a switch"),
+    cl::Hidden, cl::init(200));
+
 static cl::opt<unsigned>
     CostThreshold("dfa-cost-threshold",
                   cl::desc("Maximum cost accepted for the transformation"),
@@ -415,7 +420,7 @@ inline raw_ostream &operator<<(raw_ostream &OS, const ThreadingPath &TPath) {
 
 struct MainSwitch {
   MainSwitch(SwitchInst *SI, OptimizationRemarkEmitter *ORE) {
-    if (isPredictable(SI)) {
+    if (isCandidate(SI)) {
       Instr = SI;
     } else {
       ORE->emit([&]() {
@@ -433,83 +438,60 @@ struct MainSwitch {
   }
 
 private:
-  /// Do a use-def chain traversal. Make sure the value of the switch variable
-  /// is always a known constant. This means that all conditional jumps based on
-  /// switch variable can be converted to unconditional jumps.
-  bool isPredictable(const SwitchInst *SI) {
-    std::deque<Instruction *> Q;
+  /// Do a use-def chain traversal starting from the switch condition to see if
+  /// \p SI is a potential condidate.
+  ///
+  /// Also, collect select instructions to unfold.
+  bool isCandidate(const SwitchInst *SI) {
+    std::deque<Value *> Q;
     SmallSet<Value *, 16> SeenValues;
     SelectInsts.clear();
 
-    Value *FirstDef = SI->getOperand(0);
-    auto *Inst = dyn_cast<Instruction>(FirstDef);
-
-    // If this is a function argument or another non-instruction, then give up.
-    // We are interested in loop local variables.
-    if (!Inst)
+    Value *SICond = SI->getCondition();
+    LLVM_DEBUG(dbgs() << "\tSICond: " << *SICond << "\n");
+    if (!isa<PHINode>(SICond))
       return false;
 
-    // Require the first definition to be a PHINode
-    if (!isa<PHINode>(Inst))
-      return false;
-
-    LLVM_DEBUG(dbgs() << "\tisPredictable() FirstDef: " << *Inst << "\n");
-
-    Q.push_back(Inst);
-    SeenValues.insert(FirstDef);
+    addToQueue(SICond, Q, SeenValues);
 
     while (!Q.empty()) {
-      Instruction *Current = Q.front();
+      Value *Current = Q.front();
       Q.pop_front();
 
       if (auto *Phi = dyn_cast<PHINode>(Current)) {
         for (Value *Incoming : Phi->incoming_values()) {
-          if (!isPredictableValue(Incoming, SeenValues))
-            return false;
-          addInstToQueue(Incoming, Q, SeenValues);
+          addToQueue(Incoming, Q, SeenValues);
         }
-        LLVM_DEBUG(dbgs() << "\tisPredictable() phi: " << *Phi << "\n");
+        LLVM_DEBUG(dbgs() << "\tphi: " << *Phi << "\n");
       } else if (SelectInst *SelI = dyn_cast<SelectInst>(Current)) {
         if (!isValidSelectInst(SelI))
           return false;
-        if (!isPredictableValue(SelI->getTrueValue(), SeenValues) ||
-            !isPredictableValue(SelI->getFalseValue(), SeenValues)) {
-          return false;
-        }
-        addInstToQueue(SelI->getTrueValue(), Q, SeenValues);
-        addInstToQueue(SelI->getFalseValue(), Q, SeenValues);
-        LLVM_DEBUG(dbgs() << "\tisPredictable() select: " << *SelI << "\n");
+        addToQueue(SelI->getTrueValue(), Q, SeenValues);
+        addToQueue(SelI->getFalseValue(), Q, SeenValues);
+        LLVM_DEBUG(dbgs() << "\tselect: " << *SelI << "\n");
         if (auto *SelIUse = dyn_cast<PHINode>(SelI->user_back()))
           SelectInsts.push_back(SelectInstToUnfold(SelI, SelIUse));
+      } else if (isa<Constant>(Current)) {
+        LLVM_DEBUG(dbgs() << "\tconst: " << *Current << "\n");
+        continue;
       } else {
-        // If it is neither a phi nor a select, then we give up.
-        return false;
+        LLVM_DEBUG(dbgs() << "\tother: " << *Current << "\n");
+        // Allow unpredictable values. The hope is that those will be the
+        // initial switch values that can be ignored (they will hit the
+        // unthreaded switch) but this assumption will get checked later after
+        // paths have been enumerated (in function getStateDefMap).
+        continue;
       }
     }
 
     return true;
   }
 
-  bool isPredictableValue(Value *InpVal, SmallSet<Value *, 16> &SeenValues) {
-    if (SeenValues.contains(InpVal))
-      return true;
-
-    if (isa<ConstantInt>(InpVal))
-      return true;
-
-    // If this is a function argument or another non-instruction, then give up.
-    if (!isa<Instruction>(InpVal))
-      return false;
-
-    return true;
-  }
-
-  void addInstToQueue(Value *Val, std::deque<Instruction *> &Q,
-                      SmallSet<Value *, 16> &SeenValues) {
+  void addToQueue(Value *Val, std::deque<Value *> &Q,
+                  SmallSet<Value *, 16> &SeenValues) {
     if (SeenValues.contains(Val))
       return;
-    if (Instruction *I = dyn_cast<Instruction>(Val))
-      Q.push_back(I);
+    Q.push_back(Val);
     SeenValues.insert(Val);
   }
 
@@ -563,7 +545,16 @@ struct AllSwitchPaths {
   void run() {
     VisitedBlocks Visited;
     PathsType LoopPaths = paths(SwitchBlock, Visited, /* PathDepth = */ 1);
-    StateDefMap StateDef = getStateDefMap();
+    StateDefMap StateDef = getStateDefMap(LoopPaths);
+
+    if (StateDef.empty()) {
+      ORE->emit([&]() {
+        return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable",
+                                        Switch)
+               << "Switch instruction is not predictable.";
+      });
+      return;
+    }
 
     for (PathType Path : LoopPaths) {
       ThreadingPath TPath;
@@ -638,6 +629,9 @@ struct AllSwitchPaths {
         PathType NewPath(Path);
         NewPath.push_front(BB);
         Res.push_back(NewPath);
+        if (Res.size() >= MaxNumPaths) {
+          return Res;
+        }
       }
     }
     // This block could now be visited again from a different predecessor. Note
@@ -648,14 +642,22 @@ struct AllSwitchPaths {
   }
 
   /// Walk the use-def chain and collect all the state-defining instructions.
-  StateDefMap getStateDefMap() const {
+  ///
+  /// Return an empty map if unpredictable values encountered inside the basic
+  /// blocks of \p LoopPaths.
+  StateDefMap getStateDefMap(const PathsType &LoopPaths) const {
     StateDefMap Res;
 
+    // Basic blocks belonging to any of the loops around the switch statement.
+    SmallPtrSet<BasicBlock *, 16> LoopBBs;
+    for (const PathType &Path : LoopPaths) {
+      for (BasicBlock *BB : Path)
+        LoopBBs.insert(BB);
+    }
+
     Value *FirstDef = Switch->getOperand(0);
 
-    assert(isa<PHINode>(FirstDef) && "After select unfolding, all state "
-                                     "definitions are expected to be phi "
-                                     "nodes.");
+    assert(isa<PHINode>(FirstDef) && "The first definition must be a phi.");
 
     SmallVector<PHINode *, 8> Stack;
     Stack.push_back(dyn_cast<PHINode>(FirstDef));
@@ -667,15 +669,17 @@ struct AllSwitchPaths {
       Res[CurPhi->getParent()] = CurPhi;
       SeenValues.insert(CurPhi);
 
-      for (Value *Incoming : CurPhi->incoming_values()) {
+      for (BasicBlock *IncomingBB : CurPhi->blocks()) {
+        Value *Incoming = CurPhi->getIncomingValueForBlock(IncomingBB);
+        bool IsOutsideLoops = LoopBBs.count(IncomingBB) == 0;
         if (Incoming == FirstDef || isa<ConstantInt>(Incoming) ||
-            SeenValues.contains(Incoming)) {
+            SeenValues.contains(Incoming) || IsOutsideLoops) {
           continue;
         }
 
-        assert(isa<PHINode>(Incoming) && "After select unfolding, all state "
-                                         "definitions are expected to be phi "
-                                         "nodes.");
+        // Any unpredictable value inside the loops means we must bail out.
+        if (!isa<PHINode>(Incoming))
+          return StateDefMap();
 
         Stack.push_back(cast<PHINode>(Incoming));
       }
@@ -1279,7 +1283,7 @@ bool DFAJumpThreading::run(Function &F) {
       continue;
 
     LLVM_DEBUG(dbgs() << "\nCheck if SwitchInst in BB " << BB.getName()
-                      << " is predictable\n");
+                      << " is a candidate\n");
     MainSwitch Switch(SI, ORE);
 
     if (!Switch.getInstr())

llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll

@@ -184,3 +184,68 @@ bb49:                                             ; preds = %bb43, %bb43
 bb66:                                             ; preds = %bb59
   ret i32 0
 }
+
+; Value %init is not predictable but it's okay since it is the value initial to the switch.
+define i32 @initial.value.positive1(i32 %init) {
+; CHECK: < loop.3 case2 > [ 3, loop.3 ]
+; CHECK-NEXT: < loop.3 case2 loop.1.backedge loop.1 loop.2 > [ 1, loop.1 ]
+; CHECK-NEXT: < loop.3 case2 loop.1.backedge si.unfold.false loop.1 loop.2 > [ 4, loop.1.backedge ]
+; CHECK-NEXT: < loop.3 case3 loop.2.backedge loop.2 > [ 0, loop.2.backedge ]
+; CHECK-NEXT: < loop.3 case3 case4 loop.2.backedge loop.2 > [ 3, loop.2.backedge ]
+; CHECK-NEXT: < loop.3 case3 case4 loop.1.backedge loop.1 loop.2 > [ 1, loop.1 ]
+; CHECK-NEXT: < loop.3 case3 case4 loop.1.backedge si.unfold.false loop.1 loop.2 > [ 2, loop.1.backedge ]
+; CHECK-NEXT: < loop.3 case4 loop.2.backedge loop.2 > [ 3, loop.2.backedge ]
+; CHECK-NEXT: < loop.3 case4 loop.1.backedge loop.1 loop.2 > [ 1, loop.1 ]
+; CHECK-NEXT: < loop.3 case4 loop.1.backedge si.unfold.false loop.1 loop.2 > [ 2, loop.1.backedge ]
+entry:
+  %cmp = icmp eq i32 %init, 0
+  br label %loop.1
+
+loop.1:
+  %state.1 = phi i32 [ %init, %entry ], [ %state.1.be2, %loop.1.backedge ]
+  br label %loop.2
+
+loop.2:
+  %state.2 = phi i32 [ %state.1, %loop.1 ], [ %state.2.be, %loop.2.backedge ]
+  br label %loop.3
+
+loop.3:
+  %state = phi i32 [ %state.2, %loop.2 ], [ 3, %case2 ]
+  switch i32 %state, label %infloop.i [
+    i32 2, label %case2
+    i32 3, label %case3
+    i32 4, label %case4
+    i32 0, label %case0
+    i32 1, label %case1
+  ]
+
+case2:
+  br i1 %cmp, label %loop.3, label %loop.1.backedge
+
+case3:
+  br i1 %cmp, label %loop.2.backedge, label %case4
+
+case4:
+  br i1 %cmp, label %loop.2.backedge, label %loop.1.backedge
+
+loop.1.backedge:
+  %state.1.be = phi i32 [ 2, %case4 ], [ 4, %case2 ]
+  %state.1.be2 = select i1 %cmp, i32 1, i32 %state.1.be
+  br label %loop.1
+
+loop.2.backedge:
+  %state.2.be = phi i32 [ 3, %case4 ], [ 0, %case3 ]
+  br label %loop.2
+
+case0:
+  br label %exit
+
+case1:
+  br label %exit
+
+infloop.i:
+  br label %infloop.i
+
+exit:
+  ret i32 0
+}

llvm/test/Transforms/DFAJumpThreading/negative.ll

@@ -214,3 +214,52 @@ for.inc:
 for.end:
   ret i32 0
 }
+
+declare i32 @arbitrary_function()
+
+; Don't confuse %state.2 for the initial switch value.
+define i32 @negative6(i32 %init) {
+; REMARK: SwitchNotPredictable
+; REMARK-NEXT: negative6
+entry:
+  %cmp = icmp eq i32 %init, 0
+  br label %loop.2
+
+loop.2:
+  %state.2 = call i32 @arbitrary_function()
+  br label %loop.3
+
+loop.3:
+  %state = phi i32 [ %state.2, %loop.2 ], [ 3, %case2 ]
+  switch i32 %state, label %infloop.i [
+    i32 2, label %case2
+    i32 3, label %case3
+    i32 4, label %case4
+    i32 0, label %case0
+    i32 1, label %case1
+  ]
+
+case2:
+  br label %loop.3
+
+case3:
+  br i1 %cmp, label %loop.2.backedge, label %case4
+
+case4:
+  br label %loop.2.backedge
+
+loop.2.backedge:
+  br label %loop.2
+
+case0:
+  br label %exit
+
+case1:
+  br label %exit
+
+infloop.i:
+  br label %infloop.i
+
+exit:
+  ret i32 0
+}