[llvm] [SimplifyCFG] Fold switch over ucmp/scmp to icmp and br (PR #105636)

[Yingwei Zheng via llvm-commits]

================
@@ -7131,6 +7131,59 @@ static bool simplifySwitchOfPowersOfTwo(SwitchInst *SI, IRBuilder<> &Builder,
   return true;
 }
 
+/// Fold switch over ucmp/scmp intrinsic to br if two of the switch arms have
+/// the same destination.
+static bool simplifySwitchOfCmpIntrinsic(SwitchInst *SI,
+                                         IRBuilderBase &Builder) {
+  auto *Cmp = dyn_cast<CmpIntrinsic>(SI->getCondition());
+  if (!Cmp || !Cmp->hasOneUse() || SI->getNumCases() != 2)
+    return false;
+
+  // Find which of 1, 0 or -1 is missing.
+  SmallSet<int64_t, 3> Missing;
+  Missing.insert(1);
+  Missing.insert(0);
+  Missing.insert(-1);
+  BasicBlock *Succ = nullptr;
+  for (auto &Case : SI->cases()) {
+    std::optional<int64_t> Val = Case.getCaseValue()->getValue().trySExtValue();
+    if (!Val)
+      return false;
+    if (!Missing.erase(*Val))
+      return false;
+    if (Succ && Succ != Case.getCaseSuccessor())
+      return false;
+    Succ = Case.getCaseSuccessor();
+  }
+
+  // Determine predicate for the missing case.
+  ICmpInst::Predicate Pred;
+  assert(Missing.size() == 1 && "Should have one case left");
+  switch (*Missing.begin()) {
+  case 1:
+    Pred = ICmpInst::ICMP_UGT;
+    break;
+  case 0:
+    Pred = ICmpInst::ICMP_EQ;
+    break;
+  case -1:
+    Pred = ICmpInst::ICMP_ULT;
+    break;
+  }
+  if (Cmp->isSigned())
+    Pred = ICmpInst::getSignedPredicate(Pred);
+
+  // The dominator tree does not change, because it treats multi-edges like
+  // a single edge anyway.
+  Builder.SetInsertPoint(SI->getIterator());
+  Value *ICmp = Builder.CreateICmp(Pred, Cmp->getLHS(), Cmp->getRHS());
+  Builder.CreateCondBr(ICmp, SI->getDefaultDest(), Succ);
----------------
dtcxzyw wrote:

Branch weights should be updated.


https://github.com/llvm/llvm-project/pull/105636