llvm · fhahn · Mar 27, 2025 · May 13, 2025
diff --git a/llvm/include/llvm/CodeGen/SwitchLoweringUtils.h b/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.

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -12238,6 +12238,32 @@ void SelectionDAGBuilder::lowerWorkItem(SwitchWorkListItem W, Value *Cond,
 
         break;
       }
+      case CC_And: {
+        SDLoc dl = getCurSDLoc();
+
+        const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+        EVT VT = TLI.getValueType(DAG.getDataLayout(), I->Low->getType(), true);
+        SDValue C = DAG.getConstant(*I->Low, dl, VT);
+        SDValue Zero = DAG.getConstant(0, dl, VT);
+        SDValue CondLHS = getValue(Cond);
+        SDValue And = DAG.getNode(ISD::AND, dl, C.getValueType(), CondLHS, C);
+        auto CondD = DAG.getSetCC(dl, MVT::i1, And, Zero, ISD::SETEQ);
+        SDNodeFlags Flags;
+        SDValue BrCond =
+            DAG.getNode(ISD::BRCOND, dl, MVT::Other, getControlRoot(), CondD,
+                        DAG.getBasicBlock(I->MBB), Flags);
+
+        // Insert the false branch. Do this even if it's a fall through branch,
+        // this makes it easier to do DAG optimizations which require inverting
+        // the branch condition.
+        BrCond = DAG.getNode(ISD::BR, dl, MVT::Other, BrCond,
+                             DAG.getBasicBlock(Fallthrough));
+        addSuccessorWithProb(CurMBB, I->MBB, UnhandledProbs);
+        addSuccessorWithProb(CurMBB, Fallthrough,
+                             BranchProbability::getUnknown());
+        CurMBB->normalizeSuccProbs();
+        DAG.setRoot(BrCond);
+      }
     }
     CurMBB = Fallthrough;
   }

diff --git a/llvm/lib/CodeGen/SwitchLoweringUtils.cpp b/llvm/lib/CodeGen/SwitchLoweringUtils.cpp
@@ -362,6 +362,39 @@ void SwitchCG::SwitchLowering::findBitTestClusters(CaseClusterVector &Clusters,
     }
   }
   Clusters.resize(DstIndex);
+
+  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].Low = ConstantInt::get(SI->getContext(), NewC);
+  Clusters[ZeroIdx].Kind = CC_And;
+  Clusters.erase(Clusters.begin() + Pow2Idx);
 }
 
 bool SwitchCG::SwitchLowering::buildBitTests(CaseClusterVector &Clusters,

diff --git a/llvm/test/CodeGen/AArch64/switch-cases-to-branch-and.ll b/llvm/test/CodeGen/AArch64/switch-cases-to-branch-and.ll
@@ -0,0 +1,205 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -O3 -mtriple=arm64-apple-macosx -o - %s | FileCheck %s
+
+define i32 @switch_with_matching_dests_0_and_pow2_3_cases(i8 %v) {
+; CHECK-LABEL: switch_with_matching_dests_0_and_pow2_3_cases:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    mov w8, #100 ; =0x64
+; CHECK-NEXT:    mov w9, #223 ; =0xdf
+; CHECK-NEXT:  LBB0_1: ; %loop.header
+; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    tst w0, w9
+; CHECK-NEXT:    b.eq LBB0_4
+; CHECK-NEXT:  ; %bb.2: ; %loop.header
+; CHECK-NEXT:    ; in Loop: Header=BB0_1 Depth=1
+; CHECK-NEXT:    and w10, w0, #0xff
+; CHECK-NEXT:    cmp w10, #124
+; CHECK-NEXT:    b.eq LBB0_5
+; CHECK-NEXT:  ; %bb.3: ; %loop.latch
+; CHECK-NEXT:    ; in Loop: Header=BB0_1 Depth=1
+; CHECK-NEXT:    subs w8, w8, #1
+; CHECK-NEXT:    b.ne LBB0_1
+; CHECK-NEXT:  LBB0_4:
+; CHECK-NEXT:    mov w0, #20 ; =0x14
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  LBB0_5: ; %e2
+; CHECK-NEXT:    mov w0, #30 ; =0x1e
+; CHECK-NEXT:    ret
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  switch i8 %v, label %loop.latch [
+  i8 32, label %e1
+  i8 0, label %e1
+  i8 124, label %e2
+  ]
+
+loop.latch:
+  %iv.next = add i32 %iv, 1
+  %c = icmp eq i32 %iv.next, 100
+  br i1 %c, label %e1, label %loop.header
+
+e1:
+  ret i32 20
+
+e2:
+  ret i32 30
+}
+
+define i64 @consecutive_match_both(ptr %p, i32 %param) {
+; CHECK-LABEL: consecutive_match_both:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    mov w8, #1 ; =0x1
+; CHECK-NEXT:    mov w9, #100 ; =0x64
+; CHECK-NEXT:    mov w10, #249 ; =0xf9
+; CHECK-NEXT:    lsl w8, w8, w1
+; CHECK-NEXT:    b LBB1_2
+; CHECK-NEXT:  LBB1_1: ; %loop.latch
+; CHECK-NEXT:    ; in Loop: Header=BB1_2 Depth=1
+; CHECK-NEXT:    subs w9, w9, #1
+; CHECK-NEXT:    b.eq LBB1_5
+; CHECK-NEXT:  LBB1_2: ; %loop.header
+; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    cmp w1, #7
+; CHECK-NEXT:    b.hi LBB1_1
+; CHECK-NEXT:  ; %bb.3: ; %loop.header
+; CHECK-NEXT:    ; in Loop: Header=BB1_2 Depth=1
+; CHECK-NEXT:    tst w8, w10
+; CHECK-NEXT:    b.eq LBB1_1
+; CHECK-NEXT:  ; %bb.4: ; %e0
+; CHECK-NEXT:    mov x0, xzr
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  LBB1_5:
+; CHECK-NEXT:    mov x0, #-42 ; =0xffffffffffffffd6
+; CHECK-NEXT:    ret
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  switch i32 %param, label %loop.latch [
+  i32 7, label %e0
+  i32 6, label %e0
+  i32 5, label %e0
+  i32 4, label %e0
+  i32 3, label %e0
+  i32 0, label %e0
+  ]
+
+loop.latch:
+  %iv.next = add i32 %iv, 1
+  %ec = icmp eq i32 %iv.next, 100
+  br i1 %ec, label %e1, label %loop.header
+
+e0:
+  %m = getelementptr i8, ptr %p, i64 20
+  br label %e1
+
+e1:
+  %res = phi i64 [ 0, %e0 ], [ -42, %loop.latch ]
+  ret i64 %res
+}
+
+define i64 @consecutive_match_before(ptr %p, i32 %param) {
+; CHECK-LABEL: consecutive_match_before:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    mov w8, #1 ; =0x1
+; CHECK-NEXT:    mov w9, #100 ; =0x64
+; CHECK-NEXT:    mov w10, #25 ; =0x19
+; CHECK-NEXT:    lsl w8, w8, w1
+; CHECK-NEXT:    b LBB2_2
+; CHECK-NEXT:  LBB2_1: ; %loop.latch
+; CHECK-NEXT:    ; in Loop: Header=BB2_2 Depth=1
+; CHECK-NEXT:    subs w9, w9, #1
+; CHECK-NEXT:    b.eq LBB2_5
+; CHECK-NEXT:  LBB2_2: ; %loop.header
+; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    cmp w1, #4
+; CHECK-NEXT:    b.hi LBB2_1
+; CHECK-NEXT:  ; %bb.3: ; %loop.header
+; CHECK-NEXT:    ; in Loop: Header=BB2_2 Depth=1
+; CHECK-NEXT:    tst w8, w10
+; CHECK-NEXT:    b.eq LBB2_1
+; CHECK-NEXT:  ; %bb.4: ; %e0
+; CHECK-NEXT:    mov x0, xzr
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  LBB2_5:
+; CHECK-NEXT:    mov x0, #-42 ; =0xffffffffffffffd6
+; CHECK-NEXT:    ret
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  switch i32 %param, label %loop.latch [
+  i32 4, label %e0
+  i32 3, label %e0
+  i32 0, label %e0
+  ]
+
+loop.latch:
+  %iv.next = add i32 %iv, 1
+  %ec = icmp eq i32 %iv.next, 100
+  br i1 %ec, label %e1, label %loop.header
+
+e0:
+  %m = getelementptr i8, ptr %p, i64 20
+  br label %e1
+
+e1:
+  %res = phi i64 [ 0, %e0 ], [ -42, %loop.latch ]
+  ret i64 %res
+}
+
+define i64 @consecutive_match_after(ptr %p, i32 %param) {
+; CHECK-LABEL: consecutive_match_after:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    mov w8, #1 ; =0x1
+; CHECK-NEXT:    mov w9, #100 ; =0x64
+; CHECK-NEXT:    mov w10, #49 ; =0x31
+; CHECK-NEXT:    lsl w8, w8, w1
+; CHECK-NEXT:    b LBB3_2
+; CHECK-NEXT:  LBB3_1: ; %loop.latch
+; CHECK-NEXT:    ; in Loop: Header=BB3_2 Depth=1
+; CHECK-NEXT:    subs w9, w9, #1
+; CHECK-NEXT:    b.eq LBB3_5
+; CHECK-NEXT:  LBB3_2: ; %loop.header
+; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    cmp w1, #5
+; CHECK-NEXT:    b.hi LBB3_1
+; CHECK-NEXT:  ; %bb.3: ; %loop.header
+; CHECK-NEXT:    ; in Loop: Header=BB3_2 Depth=1
+; CHECK-NEXT:    tst w8, w10
+; CHECK-NEXT:    b.eq LBB3_1
+; CHECK-NEXT:  ; %bb.4: ; %e0
+; CHECK-NEXT:    mov x0, xzr
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  LBB3_5:
+; CHECK-NEXT:    mov x0, #-42 ; =0xffffffffffffffd6
+; CHECK-NEXT:    ret
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  switch i32 %param, label %loop.latch [
+  i32 5, label %e0
+  i32 4, label %e0
+  i32 0, label %e0
+  ]
+
+loop.latch:
+  %iv.next = add i32 %iv, 1
+  %ec = icmp eq i32 %iv.next, 100
+  br i1 %ec, label %e1, label %loop.header
+
+e0:
+  %m = getelementptr i8, ptr %p, i64 20
+  br label %e1
+
+e1:
+  %res = phi i64 [ 0, %e0 ], [ -42, %loop.latch ]
+  ret i64 %res
+}