[DAG] Add SDPatternMatch m_ZExtOrSelf/m_SExtOrSelf/m_AExtOrSelf/m_Tru…

…ncOrSelf matchers
llvm · Mar 16, 2024 · ee13aaa · ee13aaa
1 parent 84b5178
commit ee13aaa
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 0 deletions.
diff --git a/llvm/include/llvm/CodeGen/SDPatternMatch.h b/llvm/include/llvm/CodeGen/SDPatternMatch.h
@@ -725,6 +725,39 @@ inline auto m_False() {
       },
       m_Value()};
 }
+
+/// Match a zext or identity
+/// Allows to peek through optional extensions
+template <typename Opnd>
+inline Or<UnaryOpc_match<Opnd>, Opnd> m_ZExtOrSelf(Opnd &&Op) {
+  return Or<UnaryOpc_match<Opnd>, Opnd>(m_ZExt(std::forward<Opnd>(Op)),
+                                        std::forward<Opnd>(Op));
+}
+
+/// Match a sext or identity
+/// Allows to peek through optional extensions
+template <typename Opnd>
+inline Or<UnaryOpc_match<Opnd>, Opnd> m_SExtOrSelf(Opnd &&Op) {
+  return Or<UnaryOpc_match<Opnd>, Opnd>(m_SExt(std::forward<Opnd>(Op)),
+                                        std::forward<Opnd>(Op));
+}
+
+/// Match a aext or identity
+/// Allows to peek through optional extensions
+template <typename Opnd>
+inline Or<UnaryOpc_match<Opnd>, Opnd> m_AExtOrSelf(Opnd &&Op) {
+  return Or<UnaryOpc_match<Opnd>, Opnd>(m_AnyExt(std::forward<Opnd>(Op)),
+                                        std::forward<Opnd>(Op));
+}
+
+/// Match a trunc or identity
+/// Allows to peek through optional truncations
+template <typename Opnd>
+inline Or<UnaryOpc_match<Opnd>, Opnd> m_TruncOrSelf(Opnd &&Op) {
+  return Or<UnaryOpc_match<Opnd>, Opnd>(m_Trunc(std::forward<Opnd>(Op)),
+                                        std::forward<Opnd>(Op));
+}
+
 } // namespace SDPatternMatch
 } // namespace llvm
 #endif
diff --git a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
@@ -241,6 +241,29 @@ TEST_F(SelectionDAGPatternMatchTest, patternCombinators) {
   EXPECT_TRUE(sd_match(Add, m_AllOf(m_Opc(ISD::ADD), m_OneUse())));
 }
 
+TEST_F(SelectionDAGPatternMatchTest, optionalResizing) {
+  SDLoc DL;
+  auto Int32VT = EVT::getIntegerVT(Context, 32);
+  auto Int64VT = EVT::getIntegerVT(Context, 64);
+
+  SDValue Op32 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 1, Int32VT);
+  SDValue Op64 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 1, Int64VT);
+  SDValue ZExt = DAG->getNode(ISD::ZERO_EXTEND, DL, Int64VT, Op32);
+  SDValue SExt = DAG->getNode(ISD::SIGN_EXTEND, DL, Int64VT, Op32);
+  SDValue AExt = DAG->getNode(ISD::ANY_EXTEND, DL, Int64VT, Op32);
+  SDValue Trunc = DAG->getNode(ISD::TRUNCATE, DL, Int32VT, Op64);
+
+  using namespace SDPatternMatch;
+  EXPECT_TRUE(sd_match(Op32, m_ZExtOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(ZExt, m_ZExtOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(Op64, m_SExtOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(SExt, m_SExtOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(Op32, m_AExtOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(AExt, m_AExtOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(Op64, m_TruncOrSelf(m_Value())));
+  EXPECT_TRUE(sd_match(Trunc, m_TruncOrSelf(m_Value())));
+}
+
 TEST_F(SelectionDAGPatternMatchTest, matchNode) {
   SDLoc DL;
   auto Int32VT = EVT::getIntegerVT(Context, 32);