[llvm] [AArch64] Fold vector select with power-of-2 bit-test to CMTST+BSP (PR #209100)

Mugundan S via llvm-commits llvm-commits at lists.llvm.org
Sat Aug 1 08:25:41 PDT 2026


https://github.com/MGN-GIT updated https://github.com/llvm/llvm-project/pull/209100

>From 63696dd9c660a1bb1ab5b609e3335f4ecb991b33 Mon Sep 17 00:00:00 2001
From: Greenie0701 <smugundan12a at gmail.com>
Date: Mon, 13 Jul 2026 12:52:50 +0530
Subject: [PATCH 1/3] [AArch64] Fold vector select with power-of-2 bit-test to
 CMTST+BSP

---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 53 ++++++++++++++++++
 .../CodeGen/AArch64/cmtst-select-pow2-mask.ll | 55 +++++++++++++++++++
 2 files changed, 108 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/cmtst-select-pow2-mask.ll

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9872ec9a4a1fe..ab35b57b21db2 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -16529,6 +16529,59 @@ static SDValue tryLowerToBSL(SDValue N, SelectionDAG &DAG) {
                            N0->getOperand(1 - i), N1->getOperand(1 - j));
     }
 
+  // Fold: or(and(xor(setcc(and(X,Mask), Mask, eq), -1), A), and(setcc(and(X,Mask), Mask, eq), B))
+  // --> BSP(setcc(and(X,Mask), 0, ne), A, B)
+  // (X & Mask) == Mask, for a power-of-2 Mask, is equivalent to (X & Mask) != 0.
+  // The latter lowers to CMTST (one instruction) instead of AND+CMEQ (two instructions).
+  for (int i = 1; i >= 0; --i)
+    for (int j = 1; j >= 0; --j) {
+      SDValue NotMask = N0->getOperand(i);
+      SDValue A = N0->getOperand(1 - i);
+      SDValue Mask = N1->getOperand(j);
+      SDValue B = N1->getOperand(1 - j);
+
+      if (NotMask.getOpcode() != ISD::XOR ||
+          !ISD::isBuildVectorAllOnes(NotMask.getOperand(1).getNode()))
+        continue;
+      if (Mask != NotMask.getOperand(0))
+        continue;
+      if (Mask.getOpcode() != ISD::SETCC)
+        continue;
+
+      ISD::CondCode CC = cast<CondCodeSDNode>(Mask.getOperand(2))->get();
+      if (CC != ISD::SETEQ && CC != ISD::SETNE)
+        continue;
+
+      SDValue InnerAND = Mask.getOperand(0);
+      SDValue CmpRHS = Mask.getOperand(1);
+
+      if (InnerAND.getOpcode() != ISD::AND)
+        continue;
+
+      APInt SplatVal;
+      bool Op0IsPow2 = ISD::isConstantSplatVector(
+                           InnerAND.getOperand(0).getNode(), SplatVal) &&
+                       SplatVal.isPowerOf2();
+      bool Op1IsPow2 = !Op0IsPow2 &&
+                       ISD::isConstantSplatVector(
+                           InnerAND.getOperand(1).getNode(), SplatVal) &&
+                       SplatVal.isPowerOf2();
+      if (!Op0IsPow2 && !Op1IsPow2)
+        continue;
+
+      bool RHSIsZero = ISD::isBuildVectorAllZeros(CmpRHS.getNode());
+      APInt RHSSplat;
+      bool RHSIsMask = !RHSIsZero &&
+                       ISD::isConstantSplatVector(CmpRHS.getNode(), RHSSplat) &&
+                       RHSSplat == SplatVal;
+      if (!RHSIsZero && !RHSIsMask)
+        continue;
+
+      SDValue Zero = DAG.getConstant(0, DL, CmpRHS.getValueType());
+      SDValue NewMask =
+          DAG.getSetCC(DL, Mask.getValueType(), InnerAND, Zero, ISD::SETNE);
+      return DAG.getNode(AArch64ISD::BSP, DL, VT, NewMask, A, B);
+    }
   return SDValue();
 }
 
diff --git a/llvm/test/CodeGen/AArch64/cmtst-select-pow2-mask.ll b/llvm/test/CodeGen/AArch64/cmtst-select-pow2-mask.ll
new file mode 100644
index 0000000000000..60c837be638be
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/cmtst-select-pow2-mask.ll
@@ -0,0 +1,55 @@
+; Test: (X & Mask) == Mask, for a power-of-2 Mask, folds to CMTST
+; instead of AND+CMEQ when used as a select condition inside a BSL/BSP fold.
+; RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+neon < %s | FileCheck %s
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+
+define <16 x i8> @cmtst_select_v16i8_pow2(<16 x i8> %x, <16 x i8> %y) {
+; CHECK-LABEL: cmtst_select_v16i8_pow2:
+; CHECK:         movi v2.16b, #2
+; CHECK-NEXT:    cmtst v2.16b, v0.16b, v2.16b
+; CHECK-NEXT:    bif v0.16b, v1.16b, v2.16b
+; CHECK-NEXT:    ret
+  %mask = and <16 x i8> %x, splat(i8 2)
+  %cmp  = icmp eq <16 x i8> %mask, splat(i8 2)
+  %sel  = select <16 x i1> %cmp, <16 x i8> %x, <16 x i8> %y
+  ret <16 x i8> %sel
+}
+
+define <8 x i16> @cmtst_select_v8i16_pow2(<8 x i16> %x, <8 x i16> %y) {
+; CHECK-LABEL: cmtst_select_v8i16_pow2:
+; CHECK:         movi v2.8h, #4
+; CHECK-NEXT:    cmtst v2.8h, v0.8h, v2.8h
+; CHECK-NEXT:    bif v0.16b, v1.16b, v2.16b
+; CHECK-NEXT:    ret
+  %mask = and <8 x i16> %x, splat(i16 4)
+  %cmp  = icmp eq <8 x i16> %mask, splat(i16 4)
+  %sel  = select <8 x i1> %cmp, <8 x i16> %x, <8 x i16> %y
+  ret <8 x i16> %sel
+}
+
+define <4 x i32> @cmtst_select_v4i32_pow2(<4 x i32> %x, <4 x i32> %y) {
+; CHECK-LABEL: cmtst_select_v4i32_pow2:
+; CHECK:         movi v2.4s, #8
+; CHECK-NEXT:    cmtst v2.4s, v0.4s, v2.4s
+; CHECK-NEXT:    bif v0.16b, v1.16b, v2.16b
+; CHECK-NEXT:    ret
+  %mask = and <4 x i32> %x, splat(i32 8)
+  %cmp  = icmp eq <4 x i32> %mask, splat(i32 8)
+  %sel  = select <4 x i1> %cmp, <4 x i32> %x, <4 x i32> %y
+  ret <4 x i32> %sel
+}
+
+; Negative test - non-power-of-2 mask must NOT use CMTST; must fall back to
+; AND+CMEQ.
+define <16 x i8> @no_cmtst_non_pow2(<16 x i8> %x, <16 x i8> %y) {
+; CHECK-LABEL: no_cmtst_non_pow2:
+; CHECK:         movi v2.16b, #3
+; CHECK-NEXT:    and v3.16b, v0.16b, v2.16b
+; CHECK-NEXT:    cmeq v2.16b, v3.16b, v2.16b
+; CHECK-NEXT:    bif v0.16b, v1.16b, v2.16b
+; CHECK-NEXT:    ret
+  %mask = and <16 x i8> %x, splat(i8 3)
+  %cmp  = icmp eq <16 x i8> %mask, splat(i8 3)
+  %sel  = select <16 x i1> %cmp, <16 x i8> %x, <16 x i8> %y
+  ret <16 x i8> %sel
+}

>From dfd3576bb5ecdb975336fb83990999cc97b4c0c5 Mon Sep 17 00:00:00 2001
From: Mugundan S <137760120+MGN-GIT at users.noreply.github.com>
Date: Wed, 22 Jul 2026 10:27:00 +0530
Subject: [PATCH 2/3] Fix clang formatting

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index ab35b57b21db2..013b0f3af2215 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -16529,10 +16529,12 @@ static SDValue tryLowerToBSL(SDValue N, SelectionDAG &DAG) {
                            N0->getOperand(1 - i), N1->getOperand(1 - j));
     }
 
-  // Fold: or(and(xor(setcc(and(X,Mask), Mask, eq), -1), A), and(setcc(and(X,Mask), Mask, eq), B))
+  // Fold: or(and(xor(setcc(and(X,Mask), Mask, eq), -1), A),
+  // and(setcc(and(X,Mask), Mask, eq), B))
   // --> BSP(setcc(and(X,Mask), 0, ne), A, B)
-  // (X & Mask) == Mask, for a power-of-2 Mask, is equivalent to (X & Mask) != 0.
-  // The latter lowers to CMTST (one instruction) instead of AND+CMEQ (two instructions).
+  // (X & Mask) == Mask, for a power-of-2 Mask, is equivalent to (X & Mask) !=
+  // 0. The latter lowers to CMTST (one instruction) instead of AND+CMEQ (two
+  // instructions).
   for (int i = 1; i >= 0; --i)
     for (int j = 1; j >= 0; --j) {
       SDValue NotMask = N0->getOperand(i);

>From 418bdb7155966c42f93ff2b7e4828c402dffb13c Mon Sep 17 00:00:00 2001
From: Greenie0701 <smugundan12a at gmail.com>
Date: Sat, 1 Aug 2026 20:52:41 +0530
Subject: [PATCH 3/3] [AArch64] Fold vector select with power-of-2 bit-test to
 CMTST+BSP

---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 100 +++++++++---------
 llvm/lib/Target/AArch64/AArch64InstrInfo.td   |   8 +-
 2 files changed, 58 insertions(+), 50 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 013b0f3af2215..f6d90b17828c5 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -16529,61 +16529,28 @@ static SDValue tryLowerToBSL(SDValue N, SelectionDAG &DAG) {
                            N0->getOperand(1 - i), N1->getOperand(1 - j));
     }
 
-  // Fold: or(and(xor(setcc(and(X,Mask), Mask, eq), -1), A),
-  // and(setcc(and(X,Mask), Mask, eq), B))
-  // --> BSP(setcc(and(X,Mask), 0, ne), A, B)
-  // (X & Mask) == Mask, for a power-of-2 Mask, is equivalent to (X & Mask) !=
-  // 0. The latter lowers to CMTST (one instruction) instead of AND+CMEQ (two
-  // instructions).
+  // Fold: or(and(xor(AArch64ISD::CMTST(X,M), allones), A),
+  // and(AArch64ISD::CMTST(X,M), B)) --> BSP(CMTST(X,M), A, B)
+  // This absorbs the NOT produced by performSETCCCombine when it folds
+  // setcc(and(X,Mask), 0, seteq) --> NOT(CMTST(X,Mask)).
   for (int i = 1; i >= 0; --i)
     for (int j = 1; j >= 0; --j) {
-      SDValue NotMask = N0->getOperand(i);
-      SDValue A = N0->getOperand(1 - i);
-      SDValue Mask = N1->getOperand(j);
-      SDValue B = N1->getOperand(1 - j);
+      SDValue NotCMTST = N0->getOperand(i);
+      SDValue A        = N0->getOperand(1 - i);
+      SDValue CMTST    = N1->getOperand(j);
+      SDValue B        = N1->getOperand(1 - j);
 
-      if (NotMask.getOpcode() != ISD::XOR ||
-          !ISD::isBuildVectorAllOnes(NotMask.getOperand(1).getNode()))
+      if (NotCMTST.getOpcode() != ISD::XOR ||
+          !ISD::isBuildVectorAllOnes(NotCMTST.getOperand(1).getNode()))
         continue;
-      if (Mask != NotMask.getOperand(0))
+      if (NotCMTST.getOperand(0) != CMTST)
         continue;
-      if (Mask.getOpcode() != ISD::SETCC)
+      if (CMTST.getOpcode() != AArch64ISD::CMTST)
         continue;
 
-      ISD::CondCode CC = cast<CondCodeSDNode>(Mask.getOperand(2))->get();
-      if (CC != ISD::SETEQ && CC != ISD::SETNE)
-        continue;
-
-      SDValue InnerAND = Mask.getOperand(0);
-      SDValue CmpRHS = Mask.getOperand(1);
-
-      if (InnerAND.getOpcode() != ISD::AND)
-        continue;
-
-      APInt SplatVal;
-      bool Op0IsPow2 = ISD::isConstantSplatVector(
-                           InnerAND.getOperand(0).getNode(), SplatVal) &&
-                       SplatVal.isPowerOf2();
-      bool Op1IsPow2 = !Op0IsPow2 &&
-                       ISD::isConstantSplatVector(
-                           InnerAND.getOperand(1).getNode(), SplatVal) &&
-                       SplatVal.isPowerOf2();
-      if (!Op0IsPow2 && !Op1IsPow2)
-        continue;
-
-      bool RHSIsZero = ISD::isBuildVectorAllZeros(CmpRHS.getNode());
-      APInt RHSSplat;
-      bool RHSIsMask = !RHSIsZero &&
-                       ISD::isConstantSplatVector(CmpRHS.getNode(), RHSSplat) &&
-                       RHSSplat == SplatVal;
-      if (!RHSIsZero && !RHSIsMask)
-        continue;
-
-      SDValue Zero = DAG.getConstant(0, DL, CmpRHS.getValueType());
-      SDValue NewMask =
-          DAG.getSetCC(DL, Mask.getValueType(), InnerAND, Zero, ISD::SETNE);
-      return DAG.getNode(AArch64ISD::BSP, DL, VT, NewMask, A, B);
+      return DAG.getNode(AArch64ISD::BSP, DL, VT, CMTST, A, B);
     }
+ 
   return SDValue();
 }
 
@@ -29021,7 +28988,44 @@ static SDValue performSETCCCombine(SDNode *N,
       ISD::isConstantSplatVector(LHS.getNode(), SplatLHSVal) &&
       SplatLHSVal.isOne())
     return DAG.getSetCC(DL, VT, DAG.getConstant(0, DL, CmpVT), RHS, ISD::SETGE);
-
+  
+  // Fold setcc(and(X, Mask), Mask/0, eq/ne) --> [not] AArch64ISD::CMTST(X, Mask)
+  // for a power of 2 splat Mask, replacing AND+CMEQ with a single CMTST.
+  // Any NOT folds away when the result feeds a BSL/BIF/BIT select.
+  if (!DCI.isBeforeLegalize() && CmpVT.isFixedLengthVector() &&
+      (Cond == ISD::SETEQ || Cond == ISD::SETNE) &&
+      LHS.getOpcode() == ISD::AND) {
+    APInt SplatVal;
+    SDValue X, MaskOp;
+    if (ISD::isConstantSplatVector(LHS.getOperand(1).getNode(), SplatVal) &&
+        SplatVal.isPowerOf2()) {
+      X = LHS.getOperand(0);
+      MaskOp = LHS.getOperand(1);
+    } else if (ISD::isConstantSplatVector(LHS.getOperand(0).getNode(),
+                                          SplatVal) &&
+               SplatVal.isPowerOf2()) {
+      X = LHS.getOperand(1);
+      MaskOp = LHS.getOperand(0);
+    }
+    if (X.getNode()) {
+      bool RHSIsZero = ISD::isBuildVectorAllZeros(RHS.getNode());
+      APInt RHSSplat;
+      bool RHSIsMask = !RHSIsZero &&
+                       ISD::isConstantSplatVector(RHS.getNode(), RHSSplat) &&
+                       RHSSplat == SplatVal;
+      if (RHSIsZero || RHSIsMask) {
+        SDValue CMTSTNode =
+            DAG.getNode(AArch64ISD::CMTST, DL, CmpVT, X, MaskOp);
+        // CMTST gives all-ones where (X & Mask) != 0, i.e. SETNE(AND, 0).
+        // Invert when the original condition is the opposite sense.
+        bool Invert = (Cond == ISD::SETEQ) ? RHSIsZero : RHSIsMask;
+        if (Invert)
+          return DAG.getNOT(DL, CMTSTNode, CmpVT);
+        return CMTSTNode;
+      }
+    }
+  }
+ 
   return SDValue();
 }
 
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index c3ea54c1dd358..817c32e66ab38 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -950,6 +950,9 @@ def AArch64vsri : SDNode<"AArch64ISD::VSRI", SDT_AArch64vshiftinsert>;
 // element must be identical.
 def AArch64bsp: SDNode<"AArch64ISD::BSP", SDT_AArch64trivec>;
 
+// AArch64ISD::CMTST node: result is all-ones per lane where (X & Y) != 0.
+def AArch64cmtst: SDNode<"AArch64ISD::CMTST", SDT_AArch64Zip>;
+
 def AArch64cmeq : PatFrag<(ops node:$lhs, node:$rhs),
                           (setcc node:$lhs, node:$rhs, SETEQ)>;
 def AArch64cmge : PatFrag<(ops node:$lhs, node:$rhs),
@@ -978,7 +981,7 @@ def AArch64cmlez : PatFrag<(ops node:$lhs),
 def AArch64cmltz : PatFrag<(ops node:$lhs),
                            (setcc immAllZerosV, node:$lhs, SETGT)>;
 
-def AArch64cmtst : PatFrag<(ops node:$LHS, node:$RHS),
+def AArch64cmtst_frag : PatFrag<(ops node:$LHS, node:$RHS),
                            (vnot (AArch64cmeqz (and node:$LHS, node:$RHS)))>;
 
 def AArch64fcmeqz : PatFrag<(ops node:$lhs),
@@ -6109,9 +6112,10 @@ defm CMGE    : SIMDThreeSameVector<0, 0b00111, "cmge", AArch64cmge>;
 defm CMGT    : SIMDThreeSameVector<0, 0b00110, "cmgt", AArch64cmgt>;
 defm CMHI    : SIMDThreeSameVector<1, 0b00110, "cmhi", AArch64cmhi>;
 defm CMHS    : SIMDThreeSameVector<1, 0b00111, "cmhs", AArch64cmhs>;
-defm CMTST   : SIMDThreeSameVector<0, 0b10001, "cmtst", AArch64cmtst>;
+defm CMTST   : SIMDThreeSameVector<0, 0b10001, "cmtst", AArch64cmtst_frag>;
 foreach VT = [ v8i8, v16i8, v4i16, v8i16, v2i32, v4i32, v2i64 ] in {
 def : Pat<(VT (vnot (AArch64cmeqz VT:$Rn))), (!cast<Instruction>("CMTST"#VT) VT:$Rn, VT:$Rn)>;
+def : Pat<(VT (AArch64cmtst_node VT:$Rn, VT:$Rm)), (!cast<Instruction>("CMTST"#VT) VT:$Rn, VT:$Rm)>;
 }
 defm FABD    : SIMDThreeSameVectorFP<1,1,0b010,"fabd", int_aarch64_neon_fabd>;
 let Predicates = [HasNEON] in {



More information about the llvm-commits mailing list