[llvm] [SimplifyCFG] Add optimization for switches of powers of two (PR #70977)

[Yingwei Zheng via llvm-commits]

================
@@ -6839,6 +6836,64 @@ static bool ReduceSwitchRange(SwitchInst *SI, IRBuilder<> &Builder,
   return true;
 }
 
+static bool isSwitchOfPowersOfTwo(ArrayRef<APInt> Values) {
+  for (auto &Value : Values) {
+    if (!Value.isPowerOf2())
+      return false;
+  }
+
+  return true;
+}
+
+static bool simplifySwitchOfPowersOfTwo(SwitchInst *SI, IRBuilder<> &Builder,
+                                        const DataLayout &DL) {
+
+  auto *CondTy = cast<IntegerType>(SI->getCondition()->getType());
+
+  if (CondTy->getIntegerBitWidth() > 64 ||
+      !DL.fitsInLegalInteger(CondTy->getIntegerBitWidth()))
+    return false;
+
+  // Only bother with this optimization if there are more than 3 switch cases;
+  // SDAG will only bother creating jump tables for 4 or more cases.
+  if (SI->getNumCases() < 4)
+    return false;
+  SmallVector<APInt, 4> Values;
+  for (const auto &Case : SI->cases())
+    Values.push_back(Case.getCaseValue()->getValue());
+
+  // We perform this optimization only for switches with
+  // unreachable default case.
+  // This assumtion will save us from checking if `Condition` is a power of two
+  bool HasDefault =
+      !isa<UnreachableInst>(SI->getDefaultDest()->getFirstNonPHIOrDbg());
+
+  if (HasDefault || !isSwitchOfPowersOfTwo(Values))
+    return false;
+
+  Builder.SetInsertPoint(SI);
+
+  auto *Condition = SI->getCondition();
+  auto &Context = SI->getContext();
+
+  // FIXME maybe we should check if cttz intrinsic is cheap on the target
+  // architecture
+  auto *ConditionTrailingZeros =
+      Builder.CreateIntrinsic(Intrinsic::cttz, {Condition->getType()},
+                              {Condition, ConstantInt::getTrue(Context)});
+
+  SI->replaceUsesOfWith(Condition, ConditionTrailingZeros);
+
+  // Replace each case with its trailing zeros number
+  for (auto &Case : SI->cases()) {
+    auto *OrigValue = Case.getCaseValue();
+    Case.setValue(cast<ConstantInt>(ConstantInt::get(
+        OrigValue->getType(), OrigValue->getValue().countr_zero())));
----------------
dtcxzyw wrote:

```suggestion
    Case.setValue(ConstantInt::get(
        OrigValue->getType(), OrigValue->getValue().countr_zero()));
```

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