[Attributor] Simplify switches with more than one potential condition

Before, we allowed the condition to be simplified to a simple constant
only, otherwise we assumed all successors are live. Now we allow
multiple constants, and mark the default successor as dead accordingly.

Author: jdoerfert

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -4776,23 +4776,53 @@ identifyAliveSuccessors(Attributor &A, const SwitchInst &SI,
                         AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
   bool UsedAssumedInformation = false;
-  std::optional<Constant *> C =
-      A.getAssumedConstant(*SI.getCondition(), AA, UsedAssumedInformation);
-  if (!C || isa_and_nonnull<UndefValue>(*C)) {
-    // No value yet, assume all edges are dead.
-  } else if (isa_and_nonnull<ConstantInt>(*C)) {
-    for (const auto &CaseIt : SI.cases()) {
-      if (CaseIt.getCaseValue() == *C) {
-        AliveSuccessors.push_back(&CaseIt.getCaseSuccessor()->front());
-        return UsedAssumedInformation;
-      }
-    }
-    AliveSuccessors.push_back(&SI.getDefaultDest()->front());
+  SmallVector<AA::ValueAndContext> Values;
+  if (!A.getAssumedSimplifiedValues(IRPosition::value(*SI.getCondition()), &AA,
+                                    Values, AA::AnyScope,
+                                    UsedAssumedInformation)) {
+    // Something went wrong, assume all successors are live.
+    for (const BasicBlock *SuccBB : successors(SI.getParent()))
+      AliveSuccessors.push_back(&SuccBB->front());
+    return false;
+  }
+
+  if (Values.empty() ||
+      (Values.size() == 1 &&
+       isa_and_nonnull<UndefValue>(Values.front().getValue()))) {
+    // No valid value yet, assume all edges are dead.
     return UsedAssumedInformation;
-  } else {
+  }
+
+  Type &Ty = *SI.getCondition()->getType();
+  SmallPtrSet<ConstantInt *, 8> Constants;
+  auto CheckForConstantInt = [&](Value *V) {
+    if (auto *CI = dyn_cast_if_present<ConstantInt>(AA::getWithType(*V, Ty))) {
+      Constants.insert(CI);
+      return true;
+    }
+    return false;
+  };
+
+  if (!all_of(Values, [&](AA::ValueAndContext &VAC) {
+        return CheckForConstantInt(VAC.getValue());
+      })) {
     for (const BasicBlock *SuccBB : successors(SI.getParent()))
       AliveSuccessors.push_back(&SuccBB->front());
+    return UsedAssumedInformation;
   }
+
+  unsigned MatchedCases = 0;
+  for (const auto &CaseIt : SI.cases()) {
+    if (Constants.count(CaseIt.getCaseValue())) {
+      ++MatchedCases;
+      AliveSuccessors.push_back(&CaseIt.getCaseSuccessor()->front());
+    }
+  }
+
+  // If all potential values have been matched, we will not visit the default
+  // case.
+  if (MatchedCases < Constants.size())
+    AliveSuccessors.push_back(&SI.getDefaultDest()->front());
   return UsedAssumedInformation;
 }
 

llvm/test/Transforms/Attributor/value-simplify.ll

@@ -1505,6 +1505,134 @@ define i1 @constexpr_icmp2() {
   ret i1 icmp eq (ptr addrspacecast (ptr addrspace(4) @str to ptr), ptr null)
 }
 
+define i8 @switch(i1 %c1, i1 %c2) {
+; TUNIT: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; TUNIT-LABEL: define {{[^@]+}}@switch
+; TUNIT-SAME: (i1 noundef [[C1:%.*]], i1 [[C2:%.*]]) #[[ATTR2]] {
+; TUNIT-NEXT:  entry:
+; TUNIT-NEXT:    br i1 [[C1]], label [[T:%.*]], label [[F:%.*]]
+; TUNIT:       t:
+; TUNIT-NEXT:    br label [[M:%.*]]
+; TUNIT:       f:
+; TUNIT-NEXT:    br label [[M]]
+; TUNIT:       m:
+; TUNIT-NEXT:    [[J:%.*]] = phi i32 [ 0, [[T]] ], [ 4, [[F]] ]
+; TUNIT-NEXT:    switch i32 [[J]], label [[DEFAULT1:%.*]] [
+; TUNIT-NEXT:    i32 1, label [[DEAD1:%.*]]
+; TUNIT-NEXT:    i32 2, label [[DEAD2:%.*]]
+; TUNIT-NEXT:    i32 3, label [[DEAD3:%.*]]
+; TUNIT-NEXT:    i32 4, label [[ALIVE1:%.*]]
+; TUNIT-NEXT:    ]
+; TUNIT:       default1:
+; TUNIT-NEXT:    br label [[ALIVE1]]
+; TUNIT:       alive1:
+; TUNIT-NEXT:    [[K:%.*]] = phi i32 [ 1, [[M]] ], [ 4, [[DEFAULT1]] ]
+; TUNIT-NEXT:    switch i32 [[K]], label [[DEAD4:%.*]] [
+; TUNIT-NEXT:    i32 1, label [[END1:%.*]]
+; TUNIT-NEXT:    i32 2, label [[DEAD5:%.*]]
+; TUNIT-NEXT:    i32 4, label [[END2:%.*]]
+; TUNIT-NEXT:    ]
+; TUNIT:       end1:
+; TUNIT-NEXT:    ret i8 -1
+; TUNIT:       end2:
+; TUNIT-NEXT:    ret i8 -2
+; TUNIT:       dead1:
+; TUNIT-NEXT:    unreachable
+; TUNIT:       dead2:
+; TUNIT-NEXT:    unreachable
+; TUNIT:       dead3:
+; TUNIT-NEXT:    unreachable
+; TUNIT:       dead4:
+; TUNIT-NEXT:    unreachable
+; TUNIT:       dead5:
+; TUNIT-NEXT:    unreachable
+;
+; CGSCC: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; CGSCC-LABEL: define {{[^@]+}}@switch
+; CGSCC-SAME: (i1 noundef [[C1:%.*]], i1 [[C2:%.*]]) #[[ATTR1]] {
+; CGSCC-NEXT:  entry:
+; CGSCC-NEXT:    br i1 [[C1]], label [[T:%.*]], label [[F:%.*]]
+; CGSCC:       t:
+; CGSCC-NEXT:    br label [[M:%.*]]
+; CGSCC:       f:
+; CGSCC-NEXT:    br label [[M]]
+; CGSCC:       m:
+; CGSCC-NEXT:    [[J:%.*]] = phi i32 [ 0, [[T]] ], [ 4, [[F]] ]
+; CGSCC-NEXT:    switch i32 [[J]], label [[DEFAULT1:%.*]] [
+; CGSCC-NEXT:    i32 1, label [[DEAD1:%.*]]
+; CGSCC-NEXT:    i32 2, label [[DEAD2:%.*]]
+; CGSCC-NEXT:    i32 3, label [[DEAD3:%.*]]
+; CGSCC-NEXT:    i32 4, label [[ALIVE1:%.*]]
+; CGSCC-NEXT:    ]
+; CGSCC:       default1:
+; CGSCC-NEXT:    br label [[ALIVE1]]
+; CGSCC:       alive1:
+; CGSCC-NEXT:    [[K:%.*]] = phi i32 [ 1, [[M]] ], [ 4, [[DEFAULT1]] ]
+; CGSCC-NEXT:    switch i32 [[K]], label [[DEAD4:%.*]] [
+; CGSCC-NEXT:    i32 1, label [[END1:%.*]]
+; CGSCC-NEXT:    i32 2, label [[DEAD5:%.*]]
+; CGSCC-NEXT:    i32 4, label [[END2:%.*]]
+; CGSCC-NEXT:    ]
+; CGSCC:       end1:
+; CGSCC-NEXT:    ret i8 -1
+; CGSCC:       end2:
+; CGSCC-NEXT:    ret i8 -2
+; CGSCC:       dead1:
+; CGSCC-NEXT:    unreachable
+; CGSCC:       dead2:
+; CGSCC-NEXT:    unreachable
+; CGSCC:       dead3:
+; CGSCC-NEXT:    unreachable
+; CGSCC:       dead4:
+; CGSCC-NEXT:    unreachable
+; CGSCC:       dead5:
+; CGSCC-NEXT:    unreachable
+;
+entry:
+  br i1 %c1, label %t, label %f
+
+t:
+  br label %m
+
+f:
+  br label %m
+
+m:
+  %j = phi i32 [ 0, %t ], [ 4, %f ]
+  switch i32 %j, label %default1 [
+  i32 1, label %dead1
+  i32 2, label %dead2
+  i32 3, label %dead3
+  i32 4, label %alive1
+  ]
+
+default1:
+  br label %alive1
+
+alive1:
+  %k = phi i32 [ 1, %m ], [ 4, %default1 ]
+  switch i32 %k, label %dead4 [
+  i32 1, label %end1
+  i32 2, label %dead5
+  i32 4, label %end2
+  ]
+
+end1:
+  ret i8 -1
+end2:
+  ret i8 -2
+dead1:
+  ret i8 1
+dead2:
+  ret i8 2
+dead3:
+  ret i8 3
+dead4:
+  ret i8 4
+dead5:
+  ret i8 5
+}
+
 ;.
 ; TUNIT: attributes #[[ATTR0:[0-9]+]] = { nocallback nofree nosync nounwind willreturn }
 ; TUNIT: attributes #[[ATTR1]] = { memory(readwrite, argmem: none) }