From e5c799960c6fc3e3a066f8273fffd448100b454d Mon Sep 17 00:00:00 2001
From: Marc Auberer <marc.auberer@chillibits.com>
Date: Fri, 15 Mar 2024 23:22:54 +0100
Subject: [PATCH] [DAG] Add SDPatternMatch
 m_ZExtOrSelf/m_SExtOrSelf/m_AExtOrSelf/m_TruncOrSelf matchers

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

diff --git a/llvm/include/llvm/CodeGen/SDPatternMatch.h b/llvm/include/llvm/CodeGen/SDPatternMatch.h
index 96ab6b160a1c7..c0550b1e188f4 100644
--- 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 ValTy>
+inline BinaryOpc_match<UnaryOpc_match<ValTy>, ValTy, true>
+m_ZExtOrSelf(const ValTy &V) {
+  return m_Or(m_ZExt(V), V);
+}
+
+/// Match a sext or identity
+/// Allows to peek through optional extensions
+template <typename ValTy>
+inline BinaryOpc_match<UnaryOpc_match<ValTy>, ValTy, true>
+m_SExtOrSelf(const ValTy &V) {
+  return m_Or(m_SExt(V), V);
+}
+
+/// Match a aext or identity
+/// Allows to peek through optional extensions
+template <typename ValTy>
+inline BinaryOpc_match<UnaryOpc_match<ValTy>, ValTy, true>
+m_AExtOrSelf(const ValTy &V) {
+  return m_Or(m_AnyExt(V), V);
+}
+
+/// Match a trunc or identity
+/// Allows to peek through optional truncations
+template <typename ValTy>
+inline BinaryOpc_match<UnaryOpc_match<ValTy>, ValTy, true>
+m_TruncOrSelf(const ValTy &V) {
+  return m_Or(m_Trunc(V), V);
+}
+
 } // namespace SDPatternMatch
 } // namespace llvm
 #endif
diff --git a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
index 180d4306a470f..1311b91878473 100644
--- a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
+++ b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
@@ -241,6 +241,43 @@ 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);
+
+  SDValue Op0 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 1, Int32VT);
+
+  // zero extend
+  SDValue ZExt = DAG->getNode(ISD::ZERO_EXTEND, DL, Int32VT, Op0);
+  SDValue OptionalZExtLhs = DAG->getNode(ISD::OR, DL, Int32VT, ZExt, Op0);
+  SDValue OptionalZExtRhs = DAG->getNode(ISD::OR, DL, Int32VT, Op0, ZExt);
+
+  // sign extend
+  SDValue SExt = DAG->getNode(ISD::SIGN_EXTEND, DL, Int32VT, Op0);
+  SDValue OptionalSExtLhs = DAG->getNode(ISD::OR, DL, Int32VT, SExt, Op0);
+  SDValue OptionalSExtRhs = DAG->getNode(ISD::OR, DL, Int32VT, Op0, SExt);
+
+  // any extend
+  SDValue AExt = DAG->getNode(ISD::ANY_EXTEND, DL, Int32VT, Op0);
+  SDValue OptionalAExtLhs = DAG->getNode(ISD::OR, DL, Int32VT, AExt, Op0);
+  SDValue OptionalAExtRhs = DAG->getNode(ISD::OR, DL, Int32VT, Op0, AExt);
+
+  // truncate
+  SDValue Trunc = DAG->getNode(ISD::TRUNCATE, DL, Int32VT, Op0);
+  SDValue OptionalTruncLhs = DAG->getNode(ISD::OR, DL, Int32VT, Trunc, Op0);
+  SDValue OptionalTruncRhs = DAG->getNode(ISD::OR, DL, Int32VT, Op0, Trunc);
+
+  using namespace SDPatternMatch;
+  EXPECT_TRUE(sd_match(OptionalZExtLhs, m_ZExtOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalZExtRhs, m_ZExtOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalSExtLhs, m_SExtOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalSExtRhs, m_SExtOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalAExtLhs, m_AExtOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalAExtRhs, m_AExtOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalTruncLhs, m_TruncOrSelf(m_Specific(Op0))));
+  EXPECT_TRUE(sd_match(OptionalTruncRhs, m_TruncOrSelf(m_Specific(Op0))));
+}
+
 TEST_F(SelectionDAGPatternMatchTest, matchNode) {
   SDLoc DL;
   auto Int32VT = EVT::getIntegerVT(Context, 32);