[llvm] [ConstraintElim] Support adding facts from switch terminators. (PR #67061)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/67061

After 4a5bcbd5602, switch instructions can now be handled in a straight-forward manner by adding (ICMP_EQ, ConditionVal, CaseVal) for te successor blocks per case.

>From bf1a18280875cd20a97294b054cf4a085a42df73 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Thu, 21 Sep 2023 18:33:54 +0100
Subject: [PATCH] [ConstraintElim] Support adding facts from switch
 terminators.

After 4a5bcbd5602, switch instructions can now be handled in a
straight-forward manner by adding (ICMP_EQ, ConditionVal, CaseVal) for
te successor blocks per case.
---
 llvm/lib/Transforms/Scalar/ConstraintElimination.cpp | 12 ++++++++++++
 llvm/test/Transforms/ConstraintElimination/switch.ll | 12 +++---------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 3c47d36cbc0a0bc..fd6b6b553348511 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -831,6 +831,18 @@ void State::addInfoFor(BasicBlock &BB) {
     GuaranteedToExecute &= isGuaranteedToTransferExecutionToSuccessor(&I);
   }
 
+  if (auto *Switch = dyn_cast<SwitchInst>(BB.getTerminator())) {
+    for (auto &Case : Switch->cases()) {
+      BasicBlock *Succ = Case.getCaseSuccessor();
+      Value *V = Case.getCaseValue();
+      if (!canAddSuccessor(BB, Succ))
+        continue;
+      WorkList.emplace_back(FactOrCheck::getConditionFact(
+          DT.getNode(Succ), CmpInst::ICMP_EQ, Switch->getCondition(), V));
+    }
+    return;
+  }
+
   auto *Br = dyn_cast<BranchInst>(BB.getTerminator());
   if (!Br || !Br->isConditional())
     return;
diff --git a/llvm/test/Transforms/ConstraintElimination/switch.ll b/llvm/test/Transforms/ConstraintElimination/switch.ll
index 25e6d34f0250c45..8eae227f7334641 100644
--- a/llvm/test/Transforms/ConstraintElimination/switch.ll
+++ b/llvm/test/Transforms/ConstraintElimination/switch.ll
@@ -57,14 +57,10 @@ define i1 @simplify_based_on_switch(i32 %x) {
 ; CHECK-NEXT:    [[RES_1:%.*]] = xor i1 [[C_1]], [[C_2]]
 ; CHECK-NEXT:    ret i1 [[RES_1]]
 ; CHECK:       exit.2:
-; CHECK-NEXT:    [[T_1:%.*]] = icmp ult i32 [[X]], 7
-; CHECK-NEXT:    [[F_1:%.*]] = icmp ult i32 [[X]], 6
-; CHECK-NEXT:    [[RES_2:%.*]] = xor i1 [[T_1]], [[F_1]]
+; CHECK-NEXT:    [[RES_2:%.*]] = xor i1 true, false
 ; CHECK-NEXT:    ret i1 [[RES_2]]
 ; CHECK:       exit.3:
-; CHECK-NEXT:    [[T_2:%.*]] = icmp ult i32 [[X]], 11
-; CHECK-NEXT:    [[F_2:%.*]] = icmp ult i32 [[X]], 10
-; CHECK-NEXT:    [[RES_3:%.*]] = xor i1 [[T_2]], [[F_2]]
+; CHECK-NEXT:    [[RES_3:%.*]] = xor i1 true, false
 ; CHECK-NEXT:    ret i1 [[RES_3]]
 ;
 entry:
@@ -105,9 +101,7 @@ define i1 @simplify_based_on_switch_successor_branches(i32 %x) {
 ; CHECK-NEXT:    [[RES_1:%.*]] = xor i1 [[C_1]], [[C_2]]
 ; CHECK-NEXT:    ret i1 [[RES_1]]
 ; CHECK:       exit.2:
-; CHECK-NEXT:    [[T_1:%.*]] = icmp ult i32 [[X]], 7
-; CHECK-NEXT:    [[F_1:%.*]] = icmp ult i32 [[X]], 6
-; CHECK-NEXT:    [[RES_2:%.*]] = xor i1 [[T_1]], [[F_1]]
+; CHECK-NEXT:    [[RES_2:%.*]] = xor i1 true, false
 ; CHECK-NEXT:    call void @use(i1 [[RES_2]])
 ; CHECK-NEXT:    br label [[EXIT_3]]
 ; CHECK:       exit.3: