[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

719

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

[aqjune]

Remove the redundant assertion

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

 //===- SimplifyCFG.cpp - Code to perform CFG simplification ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 struct ConstantComparesGatherer {
   const DataLayout &DL;
 
   /// Value found for the switch comparison
   Value *CompValue = nullptr;
 
   /// Extra clause to be checked before the switch
   Value *Extra = nullptr;
+  /// Should Extra be frozen to be UB-safe?
+  bool ExtraNeedsFreeze = false;
 
   /// Set of integers to match in switch
   SmallVector<ConstantInt *, 8> Vals;
 
   /// Number of comparisons matched in the and/or chain
   unsigned UsedICmps = 0;
 
   /// Construct and compute the result for the comparison instruction Cond
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);
 
+    // Is there any 'select A, B, false' or 'select A, true, B'?
+    bool HasSelectFormOfAndOr = false;
+
     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);
+
+          if (!isa<BinaryOperator>(I)) {
+            assert(I->getOpcode() == Instruction::Select);
+            HasSelectFormOfAndOr = true;
+          }
           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.
       // We allow only one "Extra" case to be checked before the switch
       if (!Extra) {
         Extra = V;
         continue;
       }
       // Failed to parse a proper sequence, abort now
       CompValue = nullptr;
       break;
     }
+
+    if (Extra) {
+      // The Extra needs to be frozen if there was select because this might
+      // introduce branching on poison, which is UB.
+      // TODO: In some combination of select/and/or, this condition can be
+      // relieved.
+      ExtraNeedsFreeze = HasSelectFormOfAndOr;
+    }
   }
 };
 
 } // end anonymous namespace
 
 static void EraseTerminatorAndDCECond(Instruction *TI,
                                       MemorySSAUpdater *MSSAU = nullptr) {
   Instruction *Cond = nullptr;
bool SimplifyCFGOpt::SimplifyBranchOnICmpChain(BranchInst *BI,
 
   // Try to gather values from a chain of and/or to be turned into a switch
   ConstantComparesGatherer ConstantCompare(Cond, DL);
   // Unpack the result
   SmallVectorImpl<ConstantInt *> &Values = ConstantCompare.Vals;
   Value *CompVal = ConstantCompare.CompValue;
   unsigned UsedICmps = ConstantCompare.UsedICmps;
   Value *ExtraCase = ConstantCompare.Extra;
+  bool ExtraNeedsFreeze = ConstantCompare.ExtraNeedsFreeze;
 
   // 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.
   if (ExtraCase && Values.size() < 2)
     return false;
 
   // TODO: Preserve branch weight metadata, similarly to how
   // FoldValueComparisonIntoPredecessors preserves it.
 
   // Figure out which block is which destination.
   BasicBlock *DefaultBB = BI->getSuccessor(1);
   BasicBlock *EdgeBB = BI->getSuccessor(0);
   if (!TrueWhenEqual)
     std::swap(DefaultBB, EdgeBB);
 
   BasicBlock *BB = BI->getParent();
 
+  if (ExtraNeedsFreeze)
+    ExtraNeedsFreeze &= !isGuaranteedNotToBeUndefOrPoison(ExtraCase, Options.AC,
+                                                          BB->getTerminator());
+
   // MSAN does not like undefs as branch condition which can be introduced
   // with "explicit branch".
-  if (ExtraCase && BB->getParent()->hasFnAttribute(Attribute::SanitizeMemory))
+  if (ExtraCase && !ExtraNeedsFreeze &&
+      BB->getParent()->hasFnAttribute(Attribute::SanitizeMemory))
     return false;
 
   LLVM_DEBUG(dbgs() << "Converting 'icmp' chain with " << Values.size()
                     << " cases into SWITCH.  BB is:\n"
                     << *BB);
 
   SmallVector<DominatorTree::UpdateType, 2> Updates;
 
   // If there are any extra values that couldn't be folded into the switch
   // then we evaluate them with an explicit branch first. Split the block
   // right before the condbr to handle it.
   if (ExtraCase) {
     BasicBlock *NewBB =
         SplitBlock(BB, BI, DTU ? &DTU->getDomTree() : nullptr, /*LI=*/nullptr,
                    /*MSSAU=*/nullptr, "switch.early.test");
 
     // Remove the uncond branch added to the old block.
     Instruction *OldTI = BB->getTerminator();
     Builder.SetInsertPoint(OldTI);
+    if (ExtraNeedsFreeze)
+      ExtraCase = Builder.CreateFreeze(ExtraCase, ExtraCase->getName() + ".fr");
 
     if (TrueWhenEqual)
       Builder.CreateCondBr(ExtraCase, EdgeBB, NewBB);
     else
       Builder.CreateCondBr(ExtraCase, NewBB, EdgeBB);
 
     OldTI->eraseFromParent();
 

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:    [[COND_FR:%.*]] = freeze i1 [[COND:%.*]]
+; CHECK-NEXT:    br i1 [[COND_FR]], 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:    [[COND_FR:%.*]] = freeze i1 [[COND:%.*]]
+; CHECK-NEXT:    br i1 [[COND_FR]], 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