[llvm] [DAG] Add SDPatternMatch for VScale nodes; Use SDPatternMatch for VScale in some instances (PR #100756)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland created https://github.com/llvm/llvm-project/pull/100756

This PR contains two commits:

[DAG] Add SDPatternMatch for VScale nodes

[DAG][NFC] Use SDPatternMatch for VScale in some instances

I will commit them individually.

>From 296549260dbc6c3c0d56e7aae7022655291f9fc6 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Fri, 26 Jul 2024 07:10:04 -0700
Subject: [PATCH 1/2] [DAG] Add SDPatternMatch for VScale nodes

---
 llvm/include/llvm/CodeGen/SDPatternMatch.h              | 4 ++++
 llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp | 3 +++
 2 files changed, 7 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/SDPatternMatch.h b/llvm/include/llvm/CodeGen/SDPatternMatch.h
index a905c85f56b66..81eda1d576c9f 100644
--- a/llvm/include/llvm/CodeGen/SDPatternMatch.h
+++ b/llvm/include/llvm/CodeGen/SDPatternMatch.h
@@ -722,6 +722,10 @@ inline Or<UnaryOpc_match<Opnd>, Opnd> m_TruncOrSelf(const Opnd &Op) {
   return Or<UnaryOpc_match<Opnd>, Opnd>(m_Trunc(Op), Op);
 }
 
+template <typename Opnd> inline UnaryOpc_match<Opnd> m_VScale(const Opnd &Op) {
+  return UnaryOpc_match<Opnd>(ISD::VSCALE, Op);
+}
+
 // === Constants ===
 struct ConstantInt_match {
   APInt *BindVal;
diff --git a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
index e318f467bbf27..e1fb30fa910be 100644
--- a/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
+++ b/llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp
@@ -228,6 +228,8 @@ TEST_F(SelectionDAGPatternMatchTest, matchUnaryOp) {
   SDValue Neg = DAG->getNegative(Op0, DL, Int32VT);
   SDValue Not = DAG->getNOT(DL, Op0, Int32VT);
 
+  SDValue VScale = DAG->getVScale(DL, Int32VT, APInt::getMaxValue(32));
+
   using namespace SDPatternMatch;
   EXPECT_TRUE(sd_match(ZExt, m_UnaryOp(ISD::ZERO_EXTEND, m_Value())));
   EXPECT_TRUE(sd_match(SExt, m_SExt(m_Value())));
@@ -238,6 +240,7 @@ TEST_F(SelectionDAGPatternMatchTest, matchUnaryOp) {
   EXPECT_FALSE(sd_match(ZExt, m_Neg(m_Value())));
   EXPECT_FALSE(sd_match(Sub, m_Neg(m_Value())));
   EXPECT_FALSE(sd_match(Neg, m_Not(m_Value())));
+  EXPECT_TRUE(sd_match(VScale, m_VScale(m_Value())));
 }
 
 TEST_F(SelectionDAGPatternMatchTest, matchConstants) {

>From 9000f182f2a07dcaae7b099ef8ee16e6b39222c3 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Fri, 26 Jul 2024 07:10:27 -0700
Subject: [PATCH 2/2] [DAG][NFC] Use SDPatternMatch for VScale in some
 instances

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 060e66175d965..b35d08b327ef3 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -1086,10 +1086,9 @@ bool DAGCombiner::reassociationCanBreakAddressingModePattern(unsigned Opc,
   // (load/store (add/sub (add x, y), vscale))
   // (load/store (add/sub (add x, y), (lsl vscale, C)))
   // (load/store (add/sub (add x, y), (mul vscale, C)))
-  if ((N1.getOpcode() == ISD::VSCALE ||
-       ((N1.getOpcode() == ISD::SHL || N1.getOpcode() == ISD::MUL) &&
-        N1.getOperand(0).getOpcode() == ISD::VSCALE &&
-        isa<ConstantSDNode>(N1.getOperand(1)))) &&
+  if (sd_match(N1, m_AnyOf(m_VScale(m_Value()),
+                           m_Shl(m_VScale(m_Value()), m_ConstInt()),
+                           m_Mul(m_VScale(m_Value()), m_ConstInt()))) &&
       N1.getValueType().getFixedSizeInBits() <= 64) {
     int64_t ScalableOffset = N1.getOpcode() == ISD::VSCALE
                                  ? N1.getConstantOperandVal(0)
@@ -2975,8 +2974,7 @@ SDValue DAGCombiner::visitADD(SDNode *N) {
   }
 
   // fold a+vscale(c1)+vscale(c2) -> a+vscale(c1+c2)
-  if (N0.getOpcode() == ISD::ADD &&
-      N0.getOperand(1).getOpcode() == ISD::VSCALE &&
+  if (sd_match(N0, m_Add(m_Value(), m_VScale(m_Value()))) &&
       N1.getOpcode() == ISD::VSCALE) {
     const APInt &VS0 = N0.getOperand(1)->getConstantOperandAPInt(0);
     const APInt &VS1 = N1->getConstantOperandAPInt(0);