Add SD matchers and unit test coverage for ISD::VECTOR_SHUFFLE

llvm/include/llvm/CodeGen/SDPatternMatch.h

@@ -540,6 +540,48 @@ struct BinaryOpc_match {
   }
 };
 
+/// Matching while capturing mask
+template <typename T0, typename T1> struct SDShuffle_match {
+  T0 Op1;
+  T1 Op2;
+
+  ArrayRef<int> &CapturedMask;
+
+  // capturing mask
+  SDShuffle_match(const T0 &Op1, const T1 &Op2, ArrayRef<int> &MaskRef)
+      : Op1(Op1), Op2(Op2), CapturedMask(MaskRef) {}
+
+  template <typename MatchContext>
+  bool match(const MatchContext &Ctx, SDValue N) {
+    if (auto *I = dyn_cast<ShuffleVectorSDNode>(N)) {
+      CapturedMask = I->getMask();
+      return Op1.match(Ctx, I->getOperand(0)) &&
+             Op2.match(Ctx, I->getOperand(1));
+    }
+    return false;
+  }
+};
+
+/// Matching against a specific match
+template <typename T0, typename T1> struct SDShuffle_maskMatch {
+  T0 Op1;
+  T1 Op2;
+  ArrayRef<int> SpecificMask;
+
+  SDShuffle_maskMatch(const T0 &Op1, const T1 &Op2, ArrayRef<int> Mask)
+      : Op1(Op1), Op2(Op2), SpecificMask(Mask) {}
+
+  template <typename MatchContext>
+  bool match(const MatchContext &Ctx, SDValue N) {
+    if (auto *I = dyn_cast<ShuffleVectorSDNode>(N)) {
+      return Op1.match(Ctx, I->getOperand(0)) &&
+             Op2.match(Ctx, I->getOperand(1)) &&
+             std::equal(SpecificMask.begin(), SpecificMask.end(),
+                        I->getMask().begin(), I->getMask().end());
+    }
+    return false;
+  }
+};
 template <typename LHS_P, typename RHS_P, typename Pred_t,
           bool Commutable = false, bool ExcludeChain = false>
 struct MaxMin_match {
@@ -790,6 +832,18 @@ inline BinaryOpc_match<LHS, RHS> m_FRem(const LHS &L, const RHS &R) {
   return BinaryOpc_match<LHS, RHS>(ISD::FREM, L, R);
 }
 
+template <typename V1_t, typename V2_t>
+inline SDShuffle_match<V1_t, V2_t> m_Shuffle(const V1_t &v1, const V2_t &v2,
+                                             ArrayRef<int> &mask) {
+  return SDShuffle_match<V1_t, V2_t>(v1, v2, mask);
+}
+
+template <typename V1_t, typename V2_t>
+inline SDShuffle_maskMatch<V1_t, V2_t>
+m_ShuffleSpecificMask(const V1_t &v1, const V2_t &v2, ArrayRef<int> mask) {
+  return SDShuffle_maskMatch<V1_t, V2_t>(v1, v2, mask);
+}
+
 // === Unary operations ===
 template <typename Opnd_P, bool ExcludeChain = false> struct UnaryOpc_match {
   unsigned Opcode;

llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp

@@ -119,6 +119,33 @@ TEST_F(SelectionDAGPatternMatchTest, matchValueType) {
   EXPECT_FALSE(sd_match(Op2, m_ScalableVectorVT()));
 }
 
+TEST_F(SelectionDAGPatternMatchTest, matchVecShuffle) {
+  SDLoc DL;
+  auto Int32VT = EVT::getIntegerVT(Context, 32);
+  auto VInt32VT = EVT::getVectorVT(Context, Int32VT, 4);
+  SmallVector<int, 4> MaskData = {2, 0, 3, 1};
+  SmallVector<int, 4> otherMaskData = {1, 2, 3, 4};
+  ArrayRef<int> CapturedMask;
+
+  SDValue V0 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 1, VInt32VT);
+  SDValue V1 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 2, VInt32VT);
+  SDValue VecShuffleWithMask =
+      DAG->getVectorShuffle(VInt32VT, DL, V0, V1, MaskData);
+
+  using namespace SDPatternMatch;
+  EXPECT_TRUE(sd_match(VecShuffleWithMask,
+                       m_Shuffle(m_Value(), m_Value(), CapturedMask)));
+  EXPECT_TRUE(sd_match(VecShuffleWithMask,
+                       m_ShuffleSpecificMask(m_Value(), m_Value(), MaskData)));
+  EXPECT_TRUE(std::equal(MaskData.begin(), MaskData.end(), CapturedMask.begin(),
+                         CapturedMask.end()));
+  EXPECT_FALSE(
+      sd_match(VecShuffleWithMask,
+               m_ShuffleSpecificMask(m_Value(), m_Value(), otherMaskData)));
+  EXPECT_FALSE(std::equal(otherMaskData.begin(), otherMaskData.end(),
+                          CapturedMask.begin(), CapturedMask.end()));
+}
+
 TEST_F(SelectionDAGPatternMatchTest, matchTernaryOp) {
   SDLoc DL;
   auto Int32VT = EVT::getIntegerVT(Context, 32);