[SimplifyCFG] Update SimplifyBranchOnICmpChain to recognize select form of and/or

[aqjune]

Summary

This patch teaches SimplifyCFG::SimplifyBranchOnICmpChain to understand select form of
(x == C1 || x == C2 || ...) / (x != C1 && x != C2 && ...) and optimize them into switch if possible.
D93065 has more context about the transition, including links to the list of optimizations being updated.

[nikic]

I think this may not be correct if there is an ExtraCase. ExtraCase || X == C || ... should be fine as it is converted into if (!ExtraCase) goto; switch (X), but something like ... || X == C || ExtraCase would not be fine.

[aqjune]

In D93943#2475598, @nikic wrote:

I think this may not be correct if there is an ExtraCase. ExtraCase || X == C || ... should be fine as it is converted into if (!ExtraCase) goto; switch (X), but something like ... || X == C || ExtraCase would not be fine.

You're right. Actually, while updating test/Transforms/SimplifyCFG/switch_create.ll I realized this and maybe I wanted to think that it wasn't very important, *but* that's not the right direction. I'll update this patch.

[aqjune]

Insert freeze Extra if needed

[nikic]

In D93943#2475825, @aqjune wrote:

In D93943#2475598, @nikic wrote:

I think this may not be correct if there is an ExtraCase. ExtraCase || X == C || ... should be fine as it is converted into if (!ExtraCase) goto; switch (X), but something like ... || X == C || ExtraCase would not be fine.

You're right. Actually, while updating test/Transforms/SimplifyCFG/switch_create.ll I realized this and maybe I wanted to think that it wasn't very important, *but* that's not the right direction. I'll update this patch.

Thinking about this again, don't we need to always freeze the extra condition (unless known non-undef/poison)? While poison is only a problem when converting the select form into a branch, doesn't undef need to be frozen even in the and/or form? As the other operand may make it unconditionally true/false and thus well-defined.

In that case, I think we should either always insert the freeze, or not inserted it in the select form either (your original patch) with a FIXME for the more general problem.

[aqjune]

In D93943#2476723, @nikic wrote:

Thinking about this again, don't we need to always freeze the extra condition (unless known non-undef/poison)? While poison is only a problem when converting the select form into a branch, doesn't undef need to be frozen even in the and/or form? As the other operand may make it unconditionally true/false and thus well-defined.

In that case, I think we should either always insert the freeze, or not inserted it in the select form either (your original patch) with a FIXME for the more general problem.

I'm fine with either direction. I'll update this to uninsert freeze and leave FIXME.

[aqjune]

Uninsert freeze
I'll make a separate patch that inserts freeze to ExtraCase

[nikic]

LGTM

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

720

This assert seems a bit redundant to me (possibly leftover from previous implementation)?

[aqjune]

Remove the redundant assertion

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

private:
   }
 
   /// Given a potentially 'or'd or 'and'd together collection of icmp
   /// eq/ne/lt/gt instructions that compare a value against a constant, extract
   /// the value being compared, and stick the list constants into the Vals
   /// vector.
   /// One "Extra" case is allowed to differ from the other.
   void gather(Value *V) {
-    bool isEQ = (cast<Instruction>(V)->getOpcode() == Instruction::Or);
+    bool isEQ = match(V, m_LogicalOr(m_Value(), m_Value()));
 
     // Keep a stack (SmallVector for efficiency) for depth-first traversal
     SmallVector<Value *, 8> DFT;
     SmallPtrSet<Value *, 8> Visited;
 
     // Initialize
     Visited.insert(V);
     DFT.push_back(V);
 
     while (!DFT.empty()) {
       V = DFT.pop_back_val();
 
       if (Instruction *I = dyn_cast<Instruction>(V)) {
         // If it is a || (or && depending on isEQ), process the operands.
-        if (I->getOpcode() == (isEQ ? Instruction::Or : Instruction::And)) {
-          if (Visited.insert(I->getOperand(1)).second)
-            DFT.push_back(I->getOperand(1));
-          if (Visited.insert(I->getOperand(0)).second)
-            DFT.push_back(I->getOperand(0));
+        Value *Op0, *Op1;
+        if (isEQ ? match(I, m_LogicalOr(m_Value(Op0), m_Value(Op1)))
+                 : match(I, m_LogicalAnd(m_Value(Op0), m_Value(Op1)))) {
+          if (Visited.insert(Op1).second)
+            DFT.push_back(Op1);
+          if (Visited.insert(Op0).second)
+            DFT.push_back(Op0);
+
           continue;
         }
 
         // Try to match the current instruction
         if (matchInstruction(I, isEQ))
           // Match succeed, continue the loop
           continue;
       }
 
       // One element of the sequence of || (or &&) could not be match as a
       // comparison against the same value as the others.
bool SimplifyCFGOpt::SimplifyBranchOnICmpChain(BranchInst *BI,
   // If we didn't have a multiply compared value, fail.
   if (!CompVal)
     return false;
 
   // Avoid turning single icmps into a switch.
   if (UsedICmps <= 1)
     return false;
 
-  bool TrueWhenEqual = (Cond->getOpcode() == Instruction::Or);
+  bool TrueWhenEqual = match(Cond, m_LogicalOr(m_Value(), m_Value()));
 
   // There might be duplicate constants in the list, which the switch
   // instruction can't handle, remove them now.
   array_pod_sort(Values.begin(), Values.end(), ConstantIntSortPredicate);
   Values.erase(std::unique(Values.begin(), Values.end()), Values.end());
 
   // If Extra was used, we require at least two switch values to do the
   // transformation.  A switch with one value is just a conditional branch.

llvm/test/Transforms/SimplifyCFG/switch_create.ll

T:              ; preds = %0
   ret void
 F:              ; preds = %0
   call void @foo2( )
   ret void
 }
 
 define void @test1_select(i32 %V) {
 ; CHECK-LABEL: @test1_select(
-; CHECK-NEXT:    [[C1:%.*]] = icmp eq i32 [[V:%.*]], 4
-; CHECK-NEXT:    [[C2:%.*]] = icmp eq i32 [[V]], 17
-; CHECK-NEXT:    [[CN:%.*]] = select i1 [[C1]], i1 true, i1 [[C2]]
-; CHECK-NEXT:    br i1 [[CN]], label [[T:%.*]], label [[F:%.*]]
+; CHECK-NEXT:    switch i32 [[V:%.*]], label [[F:%.*]] [
+; CHECK-NEXT:    i32 17, label [[T:%.*]]
+; CHECK-NEXT:    i32 4, label [[T]]
+; CHECK-NEXT:    ]
 ; CHECK:       T:
 ; CHECK-NEXT:    call void @foo1()
 ; CHECK-NEXT:    ret void
 ; CHECK:       F:
 ; CHECK-NEXT:    call void @foo2()
 ; CHECK-NEXT:    ret void
 ;
   %C1 = icmp eq i32 %V, 4
T:              ; preds = %0
   ret void
 F:              ; preds = %0
   call void @foo2( )
   ret void
 }
 
 define void @test2_select(i32 %V) {
 ; CHECK-LABEL: @test2_select(
-; CHECK-NEXT:    [[C1:%.*]] = icmp ne i32 [[V:%.*]], 4
-; CHECK-NEXT:    [[C2:%.*]] = icmp ne i32 [[V]], 17
-; CHECK-NEXT:    [[CN:%.*]] = select i1 [[C1]], i1 [[C2]], i1 false
-; CHECK-NEXT:    br i1 [[CN]], label [[T:%.*]], label [[F:%.*]]
+; CHECK-NEXT:    switch i32 [[V:%.*]], label [[T:%.*]] [
+; CHECK-NEXT:    i32 17, label [[F:%.*]]
+; CHECK-NEXT:    i32 4, label [[F]]
+; CHECK-NEXT:    ]
 ; CHECK:       T:
 ; CHECK-NEXT:    call void @foo1()
 ; CHECK-NEXT:    ret void
 ; CHECK:       F:
 ; CHECK-NEXT:    call void @foo2()
 ; CHECK-NEXT:    ret void
 ;
   %C1 = icmp ne i32 %V, 4
T:
   ret i32 123
 F:
   ret i32 324
 
 }
 
 define i32 @test10_select(i32 %mode, i1 %Cond) {
 ; CHECK-LABEL: @test10_select(
-; CHECK-NEXT:  T:
-; CHECK-NEXT:    [[A:%.*]] = icmp ne i32 [[MODE:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = icmp ne i32 [[MODE]], 51
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i1 [[B]], i1 false
-; CHECK-NEXT:    [[D:%.*]] = select i1 [[C]], i1 [[COND:%.*]], i1 false
-; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[D]], i32 123, i32 324
-; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[SWITCH_EARLY_TEST:%.*]], label [[F:%.*]]
+; CHECK:       switch.early.test:
+; CHECK-NEXT:    switch i32 [[MODE:%.*]], label [[T:%.*]] [
+; CHECK-NEXT:    i32 51, label [[F]]
+; CHECK-NEXT:    i32 0, label [[F]]
+; CHECK-NEXT:    ]
+; CHECK:       T:
+; CHECK-NEXT:    [[MERGE:%.*]] = phi i32 [ 123, [[SWITCH_EARLY_TEST]] ], [ 324, [[F]] ]
+; CHECK-NEXT:    ret i32 [[MERGE]]
+; CHECK:       F:
+; CHECK-NEXT:    br label [[T]]
 ;
   %A = icmp ne i32 %mode, 0
   %B = icmp ne i32 %mode, 51
   %C = select i1 %A, i1 %B, i1 false
   %D = select i1 %C, i1 %Cond, i1 false
   br i1 %D, label %T, label %F
 T:
   ret i32 123
 F:
   ret i32 324
 
 }
 
 ; TODO: %Cond doesn't need freeze
 define i32 @test10_select_and(i32 %mode, i1 %Cond) {
 ; CHECK-LABEL: @test10_select_and(
-; CHECK-NEXT:  T:
-; CHECK-NEXT:    [[A:%.*]] = icmp ne i32 [[MODE:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = icmp ne i32 [[MODE]], 51
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i1 [[B]], i1 false
-; CHECK-NEXT:    [[D:%.*]] = and i1 [[C]], [[COND:%.*]]
-; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[D]], i32 123, i32 324
-; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[SWITCH_EARLY_TEST:%.*]], label [[F:%.*]]
+; CHECK:       switch.early.test:
+; CHECK-NEXT:    switch i32 [[MODE:%.*]], label [[T:%.*]] [
+; CHECK-NEXT:    i32 51, label [[F]]
+; CHECK-NEXT:    i32 0, label [[F]]
+; CHECK-NEXT:    ]
+; CHECK:       T:
+; CHECK-NEXT:    [[MERGE:%.*]] = phi i32 [ 123, [[SWITCH_EARLY_TEST]] ], [ 324, [[F]] ]
+; CHECK-NEXT:    ret i32 [[MERGE]]
+; CHECK:       F:
+; CHECK-NEXT:    br label [[T]]
 ;
   %A = icmp ne i32 %mode, 0
   %B = icmp ne i32 %mode, 51
   %C = select i1 %A, i1 %B, i1 false
   %D = and i1 %C, %Cond
   br i1 %D, label %T, label %F
 T:
   ret i32 123
 F:
   ret i32 324
 
 }
 
 define i32 @test10_select_nofreeze(i32 %mode, i1 noundef %Cond) {
 ; CHECK-LABEL: @test10_select_nofreeze(
-; CHECK-NEXT:  T:
-; CHECK-NEXT:    [[A:%.*]] = icmp ne i32 [[MODE:%.*]], 0
-; CHECK-NEXT:    [[B:%.*]] = icmp ne i32 [[MODE]], 51
-; CHECK-NEXT:    [[C:%.*]] = select i1 [[A]], i1 [[B]], i1 false
-; CHECK-NEXT:    [[D:%.*]] = select i1 [[C]], i1 [[COND:%.*]], i1 false
-; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[D]], i32 123, i32 324
-; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[SWITCH_EARLY_TEST:%.*]], label [[F:%.*]]
+; CHECK:       switch.early.test:
+; CHECK-NEXT:    switch i32 [[MODE:%.*]], label [[T:%.*]] [
+; CHECK-NEXT:    i32 51, label [[F]]
+; CHECK-NEXT:    i32 0, label [[F]]
+; CHECK-NEXT:    ]
+; CHECK:       T:
+; CHECK-NEXT:    [[MERGE:%.*]] = phi i32 [ 123, [[SWITCH_EARLY_TEST]] ], [ 324, [[F]] ]
+; CHECK-NEXT:    ret i32 [[MERGE]]
+; CHECK:       F:
+; CHECK-NEXT:    br label [[T]]
 ;
   %A = icmp ne i32 %mode, 0
   %B = icmp ne i32 %mode, 51
   %C = select i1 %A, i1 %B, i1 false
   %D = select i1 %C, i1 %Cond, i1 false
   br i1 %D, label %T, label %F
 T:
   ret i32 123
if:
   ret void
 
 else:
   ret void
 }
 
 define void @test19_select(i32 %arg) {
 ; CHECK-LABEL: @test19_select(
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[ARG:%.*]], -2
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i32 [[AND]], 12
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[ARG]], 32
-; CHECK-NEXT:    [[PRED:%.*]] = select i1 [[CMP1]], i1 true, i1 [[CMP2]]
-; CHECK-NEXT:    br i1 [[PRED]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK-NEXT:    switch i32 [[ARG:%.*]], label [[ELSE:%.*]] [
+; CHECK-NEXT:    i32 32, label [[IF:%.*]]
+; CHECK-NEXT:    i32 13, label [[IF]]
+; CHECK-NEXT:    i32 12, label [[IF]]
+; CHECK-NEXT:    ]
 ; CHECK:       if:
 ; CHECK-NEXT:    call void @foo1()
 ; CHECK-NEXT:    ret void
 ; CHECK:       else:
 ; CHECK-NEXT:    ret void
 ;
   %and = and i32 %arg, -2
   %cmp1 = icmp eq i32 %and, 12