[SwitchLowering] Support merging 0 and power-of-2 case.

Author: fhahn

llvm/include/llvm/CodeGen/GlobalISel/IRTranslator.h

@@ -405,13 +405,13 @@ class IRTranslator : public MachineFunctionPass {
       BranchProbability UnhandledProbs, SwitchCG::CaseClusterIt I,
       MachineBasicBlock *Fallthrough, bool FallthroughUnreachable);
 
-  bool lowerSwitchRangeWorkItem(SwitchCG::CaseClusterIt I, Value *Cond,
-                                MachineBasicBlock *Fallthrough,
-                                bool FallthroughUnreachable,
-                                BranchProbability UnhandledProbs,
-                                MachineBasicBlock *CurMBB,
-                                MachineIRBuilder &MIB,
-                                MachineBasicBlock *SwitchMBB);
+  bool lowerSwitchAndOrRangeWorkItem(SwitchCG::CaseClusterIt I, Value *Cond,
+                                     MachineBasicBlock *Fallthrough,
+                                     bool FallthroughUnreachable,
+                                     BranchProbability UnhandledProbs,
+                                     MachineBasicBlock *CurMBB,
+                                     MachineIRBuilder &MIB,
+                                     MachineBasicBlock *SwitchMBB);
 
   bool lowerBitTestWorkItem(
       SwitchCG::SwitchWorkListItem W, MachineBasicBlock *SwitchMBB,

llvm/include/llvm/CodeGen/SwitchLoweringUtils.h

@@ -35,7 +35,8 @@ enum CaseClusterKind {
   /// A cluster of cases suitable for jump table lowering.
   CC_JumpTable,
   /// A cluster of cases suitable for bit test lowering.
-  CC_BitTests
+  CC_BitTests,
+  CC_And
 };
 
 /// A cluster of case labels.
@@ -141,6 +142,8 @@ struct CaseBlock {
   BranchProbability TrueProb, FalseProb;
   bool IsUnpredictable;
 
+  bool EmitAnd = false;
+
   // Constructor for SelectionDAG.
   CaseBlock(ISD::CondCode cc, const Value *cmplhs, const Value *cmprhs,
             const Value *cmpmiddle, MachineBasicBlock *truebb,

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -1058,18 +1058,15 @@ bool IRTranslator::lowerJumpTableWorkItem(SwitchCG::SwitchWorkListItem W,
   }
   return true;
 }
-bool IRTranslator::lowerSwitchRangeWorkItem(SwitchCG::CaseClusterIt I,
-                                            Value *Cond,
-                                            MachineBasicBlock *Fallthrough,
-                                            bool FallthroughUnreachable,
-                                            BranchProbability UnhandledProbs,
-                                            MachineBasicBlock *CurMBB,
-                                            MachineIRBuilder &MIB,
-                                            MachineBasicBlock *SwitchMBB) {
+bool IRTranslator::lowerSwitchAndOrRangeWorkItem(
+    SwitchCG::CaseClusterIt I, Value *Cond, MachineBasicBlock *Fallthrough,
+    bool FallthroughUnreachable, BranchProbability UnhandledProbs,
+    MachineBasicBlock *CurMBB, MachineIRBuilder &MIB,
+    MachineBasicBlock *SwitchMBB) {
   using namespace SwitchCG;
   const Value *RHS, *LHS, *MHS;
   CmpInst::Predicate Pred;
-  if (I->Low == I->High) {
+  if (I->Low == I->High || I->Kind == CC_And) {
     // Check Cond == I->Low.
     Pred = CmpInst::ICMP_EQ;
     LHS = Cond;
@@ -1087,6 +1084,7 @@ bool IRTranslator::lowerSwitchRangeWorkItem(SwitchCG::CaseClusterIt I,
   // The false probability is the sum of all unhandled cases.
   CaseBlock CB(Pred, FallthroughUnreachable, LHS, RHS, MHS, I->MBB, Fallthrough,
                CurMBB, MIB.getDebugLoc(), I->Prob, UnhandledProbs);
+  CB.EmitAnd = I->Kind == CC_And;
 
   emitSwitchCase(CB, SwitchMBB, MIB);
   return true;
@@ -1326,10 +1324,11 @@ bool IRTranslator::lowerSwitchWorkItem(SwitchCG::SwitchWorkListItem W,
       }
       break;
     }
+    case CC_And:
     case CC_Range: {
-      if (!lowerSwitchRangeWorkItem(I, Cond, Fallthrough,
-                                    FallthroughUnreachable, UnhandledProbs,
-                                    CurMBB, MIB, SwitchMBB)) {
+      if (!lowerSwitchAndOrRangeWorkItem(I, Cond, Fallthrough,
+                                         FallthroughUnreachable, UnhandledProbs,
+                                         CurMBB, MIB, SwitchMBB)) {
         LLVM_DEBUG(dbgs() << "Failed to lower switch range");
         return false;
       }

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -2887,7 +2887,17 @@ void SelectionDAGBuilder::visitSwitchCase(CaseBlock &CB,
   EVT MemVT = TLI.getMemValueType(DAG.getDataLayout(), CB.CmpLHS->getType());
 
   // Build the setcc now.
-  if (!CB.CmpMHS) {
+  if (CB.EmitAnd) {
+    SDLoc dl = getCurSDLoc();
+
+    const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+    EVT VT = TLI.getValueType(DAG.getDataLayout(), CB.CmpRHS->getType(), true);
+    SDValue C = DAG.getConstant(*cast<ConstantInt>(CB.CmpRHS), dl, VT);
+    SDValue Zero = DAG.getConstant(0, dl, VT);
+    SDValue CondLHS = getValue(CB.CmpLHS);
+    SDValue And = DAG.getNode(ISD::AND, dl, C.getValueType(), CondLHS, C);
+    Cond = DAG.getSetCC(dl, MVT::i1, And, Zero, ISD::SETEQ);
+  } else if (!CB.CmpMHS) {
     // Fold "(X == true)" to X and "(X == false)" to !X to
     // handle common cases produced by branch lowering.
     if (CB.CmpRHS == ConstantInt::getTrue(*DAG.getContext()) &&
@@ -12308,10 +12318,11 @@ void SelectionDAGBuilder::lowerWorkItem(SwitchWorkListItem W, Value *Cond,
         }
         break;
       }
+      case CC_And:
       case CC_Range: {
         const Value *RHS, *LHS, *MHS;
         ISD::CondCode CC;
-        if (I->Low == I->High) {
+        if (I->Low == I->High || I->Kind == CC_And) {
           // Check Cond == I->Low.
           CC = ISD::SETEQ;
           LHS = Cond;
@@ -12333,6 +12344,7 @@ void SelectionDAGBuilder::lowerWorkItem(SwitchWorkListItem W, Value *Cond,
         CaseBlock CB(CC, LHS, RHS, MHS, I->MBB, Fallthrough, CurMBB,
                      getCurSDLoc(), I->Prob, UnhandledProbs);
 
+        CB.EmitAnd = I->Kind == CC_And;
         if (CurMBB == SwitchMBB)
           visitSwitchCase(CB, SwitchMBB);
         else

llvm/lib/CodeGen/SwitchLoweringUtils.cpp

@@ -362,6 +362,41 @@ void SwitchCG::SwitchLowering::findBitTestClusters(CaseClusterVector &Clusters,
     }
   }
   Clusters.resize(DstIndex);
+
+  // Check if the clusters contain one checking for 0 and another one checking
+  // for a power-of-2 constant with matching destinations. Those clusters can be
+  // combined to a single ane with CC_And.
+  unsigned ZeroIdx = -1;
+  for (const auto &[Idx, C] : enumerate(Clusters)) {
+    if (C.Kind != CC_Range || C.Low != C.High)
+      continue;
+    if (C.Low->isZero()) {
+      ZeroIdx = Idx;
+      break;
+    }
+  }
+  if (ZeroIdx == -1u)
+    return;
+
+  unsigned Pow2Idx = -1;
+  for (const auto &[Idx, C] : enumerate(Clusters)) {
+    if (C.Kind != CC_Range || C.Low != C.High || C.MBB != Clusters[ZeroIdx].MBB)
+      continue;
+    if (C.Low->getValue().isPowerOf2()) {
+      Pow2Idx = Idx;
+      break;
+    }
+  }
+  if (Pow2Idx == -1u)
+    return;
+
+  APInt Pow2 = Clusters[Pow2Idx].Low->getValue();
+  APInt NewC = (Pow2 + 1) * -1;
+  Clusters[ZeroIdx].Low = ConstantInt::get(SI->getContext(), NewC);
+  Clusters[ZeroIdx].High = ConstantInt::get(SI->getContext(), NewC);
+  Clusters[ZeroIdx].Kind = CC_And;
+  Clusters[ZeroIdx].Prob += Clusters[Pow2Idx].Prob;
+  Clusters.erase(Clusters.begin() + Pow2Idx);
 }
 
 bool SwitchCG::SwitchLowering::buildBitTests(CaseClusterVector &Clusters,