[llvm] [PatternMatch] match splats from an all-zeroes shufflevector (PR #228016)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Graham Hunter (huntergr-arm)

<details>
<summary>Changes</summary>

Take 2, after I mismerged the previous version.

---
Full diff: https://github.com/llvm/llvm-project/pull/228016.diff


3 Files Affected:

- (modified) llvm/include/llvm/IR/PatternMatch.h (+19-5) 
- (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+1-1) 
- (modified) llvm/unittests/IR/PatternMatch.cpp (+11) 


``````````diff
diff --git a/llvm/include/llvm/IR/PatternMatch.h b/llvm/include/llvm/IR/PatternMatch.h
index 32d5e45320463..a279d65cd15bb 100644
--- a/llvm/include/llvm/IR/PatternMatch.h
+++ b/llvm/include/llvm/IR/PatternMatch.h
@@ -218,14 +218,28 @@ template <typename SubPattern_t> struct Splat_match {
       auto *Splat = C->getSplatValue();
       return Splat ? SubPattern.match(Splat) : false;
     }
-    // TODO: Extend to other cases (e.g. shufflevectors).
-    return false;
+
+    auto *Shuffle = dyn_cast<ShuffleVectorInst>(V);
+    if (!Shuffle || !Shuffle->isZeroEltSplat())
+      return false;
+
+    // Look for an insertelement.
+    auto *Insert = dyn_cast<InsertElementInst>(Shuffle->getOperand(0));
+    if (!Insert)
+      return false;
+
+    Value *SplatElt = Insert->getOperand(1);
+    ConstantInt *Idx = dyn_cast<ConstantInt>(Insert->getOperand(2));
+    if (!Idx || !Idx->isZero())
+      return false;
+
+    return SubPattern.match(SplatElt);
   }
 };
 
-/// Match a constant splat. TODO: Extend this to non-constant splats.
-template <typename T>
-inline Splat_match<T> m_ConstantSplat(const T &SubPattern) {
+/// Match a vector splat. May be a constant splat or a shufflevector of the
+/// first element.
+template <typename T> inline Splat_match<T> m_Splat(const T &SubPattern) {
   return SubPattern;
 }
 
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 72adf311850a3..140953ccdcd94 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2470,7 +2470,7 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
           bool SplatLHS) -> Instruction * {
     Value *Idx;
     Constant *Splat, *SubVector, *Dest;
-    if (!match(MaybeSplat, m_ConstantSplat(m_Constant(Splat))) ||
+    if (!match(MaybeSplat, m_Splat(m_Constant(Splat))) ||
         !match(MaybeSubVector,
                m_VectorInsert(m_Constant(Dest), m_Constant(SubVector),
                               m_Value(Idx))))
diff --git a/llvm/unittests/IR/PatternMatch.cpp b/llvm/unittests/IR/PatternMatch.cpp
index 64ccd1fdc47dd..7bf92f7761a51 100644
--- a/llvm/unittests/IR/PatternMatch.cpp
+++ b/llvm/unittests/IR/PatternMatch.cpp
@@ -1532,6 +1532,17 @@ TEST_F(PatternMatchTest, VectorOps) {
       SP2, m_Shuffle(m_InsertElt(m_Undef(), m_Value(A), m_Zero()),
                      m_Undef(), m_ZeroMask())));
   EXPECT_TRUE(A == Val);
+
+  // Repeat the above with m_Splat.
+  EXPECT_TRUE(match(SI1, m_Splat(m_SpecificInt(1))));
+  EXPECT_FALSE(match(SI2, m_Splat(m_Value())));
+  EXPECT_FALSE(match(SI3, m_Splat(m_Value())));
+  EXPECT_FALSE(match(SI4, m_Splat(m_Value())));
+
+  A = nullptr;
+  EXPECT_TRUE(match(SP1, m_Splat(m_SpecificInt(2))));
+  EXPECT_TRUE(match(SP2, m_Splat(m_Value(A))));
+  EXPECT_TRUE(A == Val);
 }
 
 TEST_F(PatternMatchTest, UndefPoisonMix) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/228016