[llvm] Add matchers m_ExtractSubvector and m_InsertSubvector (PR #120212)

[Thor Preimesberger via llvm-commits]

https://github.com/cheezeburglar created https://github.com/llvm/llvm-project/pull/120212

For #118846

>From 24329cfc604b81c957c56216a84a599b998cf327 Mon Sep 17 00:00:00 2001
From: 4444-thor <ThorP at protonmail.com>
Date: Tue, 17 Dec 2024 02:03:28 -0600
Subject: [PATCH] Add matchers m_ExtractSubvector and m_InsertSubvector

---
 llvm/include/llvm/CodeGen/SDPatternMatch.h    | 12 +++++++++
 .../CodeGen/SelectionDAGPatternMatchTest.cpp  | 25 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/SDPatternMatch.h b/llvm/include/llvm/CodeGen/SDPatternMatch.h
index d21cc962da46cb..153f522cfb7627 100644
--- a/llvm/include/llvm/CodeGen/SDPatternMatch.h
+++ b/llvm/include/llvm/CodeGen/SDPatternMatch.h
@@ -514,6 +514,12 @@ m_InsertElt(const T0_P &Vec, const T1_P &Val, const T2_P &Idx) {
                                             Idx);
 }
 
+template <typename LHS, typename RHS, typename IDX>
+inline TernaryOpc_match<LHS, RHS, IDX>
+m_InsertSubvector(const LHS &L, const RHS &R, const IDX &I) {
+  return TernaryOpc_match<LHS, RHS, IDX>(ISD::INSERT_SUBVECTOR, L, R, I);
+}
+
 // === Binary operations ===
 template <typename LHS_P, typename RHS_P, bool Commutable = false,
           bool ExcludeChain = false>
@@ -802,6 +808,12 @@ inline BinaryOpc_match<LHS, RHS> m_ExtractElt(const LHS &Vec, const RHS &Idx) {
   return BinaryOpc_match<LHS, RHS>(ISD::EXTRACT_VECTOR_ELT, Vec, Idx);
 }
 
+template <typename LHS, typename RHS>
+inline BinaryOpc_match<LHS, RHS> m_ExtractSubvector(const LHS &L,
+                                                    const RHS &R) {
+  return BinaryOpc_match<LHS, RHS>(ISD::EXTRACT_SUBVECTOR, L, R);
+}
+
 // === Unary operations ===
 template <typename Opnd_P, bool ExcludeChain = false> struct UnaryOpc_match {
   unsigned Opcode;
diff --git a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
index 259bdad0ab2723..8cc818aca0a036 100644
--- a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
+++ b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
@@ -138,9 +138,14 @@ TEST_F(SelectionDAGPatternMatchTest, matchTernaryOp) {
   SDValue Select = DAG->getSelect(DL, MVT::i1, Cond, T, F);
 
   auto VInt32VT = EVT::getVectorVT(Context, Int32VT, 4);
+  auto SmallVInt32VT = EVT::getVectorVT(Context, Int32VT, 2);
+  auto Idx0 = DAG->getVectorIdxConstant(0, DL);
   SDValue V1 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 6, VInt32VT);
   SDValue V2 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 7, VInt32VT);
+  SDValue V3 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 8, SmallVInt32VT);
   SDValue VSelect = DAG->getNode(ISD::VSELECT, DL, VInt32VT, Cond, V1, V2);
+  SDValue InsertSubvector =
+      DAG->getNode(ISD::INSERT_SUBVECTOR, DL, VInt32VT, V2, V3, Idx0);
 
   SDValue ExtractELT =
       DAG->getNode(ISD::EXTRACT_VECTOR_ELT, DL, Int32VT, V1, Op3);
@@ -182,15 +187,23 @@ TEST_F(SelectionDAGPatternMatchTest, matchTernaryOp) {
   EXPECT_TRUE(sd_match(ExtractELT, m_ExtractElt(m_Value(), m_Value())));
   EXPECT_TRUE(sd_match(ExtractELT, m_ExtractElt(m_Value(), m_ConstInt())));
   EXPECT_TRUE(sd_match(ExtractELT, m_ExtractElt(m_Value(), m_SpecificInt(1))));
+
+  EXPECT_TRUE(sd_match(InsertSubvector,
+                       m_InsertSubvector(m_Value(), m_Value(), m_Value())));
+  EXPECT_TRUE(sd_match(
+      InsertSubvector,
+      m_InsertSubvector(m_Specific(V2), m_Specific(V3), m_Specific(Idx0))));
 }
 
 TEST_F(SelectionDAGPatternMatchTest, matchBinaryOp) {
   SDLoc DL;
   auto Int32VT = EVT::getIntegerVT(Context, 32);
   auto Float32VT = EVT::getFloatingPointVT(32);
+  auto BigVInt32VT = EVT::getVectorVT(Context, Int32VT, 8);
   auto VInt32VT = EVT::getVectorVT(Context, Int32VT, 4);
 
   SDValue V1 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 6, VInt32VT);
+  auto Idx0 = DAG->getVectorIdxConstant(0, DL);
 
   SDValue Op0 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 1, Int32VT);
   SDValue Op1 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 2, Int32VT);
@@ -233,6 +246,10 @@ TEST_F(SelectionDAGPatternMatchTest, matchBinaryOp) {
   SDValue SFAdd = DAG->getNode(ISD::STRICT_FADD, DL, {Float32VT, MVT::Other},
                                {DAG->getEntryNode(), Op2, Op2});
 
+  SDValue Vec = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 9, BigVInt32VT);
+  SDValue SubVec =
+      DAG->getNode(ISD::EXTRACT_SUBVECTOR, DL, VInt32VT, Vec, Idx0);
+
   SDValue InsertELT =
       DAG->getNode(ISD::INSERT_VECTOR_ELT, DL, VInt32VT, V1, Op0, Op4);
 
@@ -293,6 +310,14 @@ TEST_F(SelectionDAGPatternMatchTest, matchBinaryOp) {
   EXPECT_FALSE(sd_match(SFAdd, m_ChainedBinOp(ISD::STRICT_FADD, m_OtherVT(),
                                               m_SpecificVT(Float32VT))));
 
+  EXPECT_TRUE(sd_match(SubVec, m_ExtractSubvector(m_Value(), m_Value())));
+  EXPECT_TRUE(
+      sd_match(SubVec, m_BinOp(ISD::EXTRACT_SUBVECTOR, m_Value(), m_Value())));
+  EXPECT_TRUE(
+      sd_match(SubVec, m_ExtractSubvector(m_Specific(Vec), m_Specific(Idx0))));
+  EXPECT_TRUE(sd_match(SubVec, m_BinOp(ISD::EXTRACT_SUBVECTOR, m_Specific(Vec),
+                                       m_Specific(Idx0))));
+
   EXPECT_TRUE(
       sd_match(InsertELT, m_InsertElt(m_Value(), m_Value(), m_Value())));
   EXPECT_TRUE(