Update m_Undef to match vectors/aggrs with undefs and poisons mixed

[aqjune]

Summary

This is an alternative approach to fixing https://reviews.llvm.org/D93990#2666922
by teaching m_Undef to match vectors/aggrs with poison elements instead of hard-coding
it at SimplifyQuery::isUndefValue (D99853).

This patch also contains fixes in InstCombine files to use the m_Undef matcher instead
of isa<UndefValue>.

[aqjune]

llvm/test/Transforms/PhaseOrdering/X86/scalarization.ll

30 ↗

(On Diff #336143)

Interestingly, this updates caused unused lanes to be undef even if they were updated to be poison in the past.
I guess this is related to the (still) ongoing work about shufflevector's undef mask things.
The work is currently blocked by the incorrect select -> and/or transformation.
If the shufflevector patches are finally landed, this will be resolved.

[aqjune]

ping

[lebedev.ri]

llvm/include/llvm/IR/PatternMatch.h

92–100

I'm not sure i like recursion here.
How about a worklist algorithm?

[aqjune]

Use Worklist

[aqjune]

llvm/include/llvm/IR/PatternMatch.h

92–100

Thank you for the suggestion! I updated this patch.

[lebedev.ri]

llvm/include/llvm/IR/PatternMatch.h

92–100

The point of that lambda was to check non-aggregate operands on the spot, without adding them into the worklist.
If we don't want that, just inline the lambda into the loop,
and do plain Worklist.emplace_back(V); before the loop..

[aqjune]

Avoid adding elem to Worklist, clang-format

[aqjune]

llvm/include/llvm/IR/PatternMatch.h

92–100

Would this new version address the concern, maybe?

[lebedev.ri]

last few nits

llvm/include/llvm/IR/PatternMatch.h

111–112

no else after return

131–132

I think this should state what happens about partial cases,
i.e. is { undef, 0 } an undef? I guess not, it should not match.

[aqjune]

Clarify when an aggregate with non-undef/poison elements is given, address comments, add more unit tests

[lebedev.ri]

llvm/include/llvm/IR/PatternMatch.h

103

Sorry for being nitpicky, but is that what it returns?
Wouldn't it have to then return false for e.g. constant expressions, since they could be undef?

[aqjune]

Be more explicit about what CheckValue returns

[lebedev.ri]

LGTM.
Please split this patch into two - the matcher+unittest, and the usage of the matcher.

llvm/include/llvm/IR/PatternMatch.h

121–122

Hm, i'd invert the CheckValue()'s return value, this seems unnatural.

llvm/include/llvm/Analysis/InstructionSimplify.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_INSTRUCTIONSIMPLIFY_H
 #define LLVM_ANALYSIS_INSTRUCTIONSIMPLIFY_H
 
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Operator.h"
+#include "llvm/IR/PatternMatch.h"
 
 namespace llvm {
 
 template <typename T, typename... TArgs> class AnalysisManager;
 template <class T> class ArrayRef;
 class AssumptionCache;
 class BinaryOperator;
 class CallBase;
SimplifyQuery getWithoutUndef() const {
     return Copy;
   }
 
   /// If CanUseUndef is true, returns whether \p V is undef.
   /// Otherwise always return false.
   bool isUndefValue(Value *V) const {
     if (!CanUseUndef)
       return false;
-    return isa<UndefValue>(V);
+
+    using namespace PatternMatch;
+    return match(V, m_Undef());
   }
 };
 
 // NOTE: the explicit multiple argument versions of these functions are
 // deprecated.
 // Please use the SimplifyQuery versions in new code.
 
 /// Given operand for an FNeg, fold the result or return null.

llvm/include/llvm/IR/PatternMatch.h

 /// Match an arbitrary binary operation and ignore it.
 inline class_match<BinaryOperator> m_BinOp() {
   return class_match<BinaryOperator>();
 }
 
 /// Matches any compare instruction and ignore it.
 inline class_match<CmpInst> m_Cmp() { return class_match<CmpInst>(); }
 
-/// Match an arbitrary undef constant.
-inline class_match<UndefValue> m_Undef() { return class_match<UndefValue>(); }
+struct undef_match {
+  static bool check(const Value *V) {
+    if (isa<UndefValue>(V))
+      return true;
+
+    const auto *CA = dyn_cast<ConstantAggregate>(V);
+    if (!CA)
+      return false;
+
+    SmallPtrSet<const ConstantAggregate *, 8> Seen;
+    SmallVector<const ConstantAggregate *, 8> Worklist;
+
+    // Either UndefValue, PoisonValue, or an aggregate that only contains
+    // these is accepted by matcher.
+    // CheckValue returns false if CA cannot satisfy this constraint.
+    auto CheckValue = [&](const ConstantAggregate *CA) {
+      for (const Value *Op : CA->operand_values()) {
+        if (isa<UndefValue>(Op))
+          continue;
+
+        const auto *CA = dyn_cast<ConstantAggregate>(Op);
+        if (!CA)
+          return false;
+        if (Seen.insert(CA).second)
+          Worklist.emplace_back(CA);
+      }
+
+      return true;
+    };
+
+    if (!CheckValue(CA))
+      return false;
+
+    while (!Worklist.empty()) {
+      if (!CheckValue(Worklist.pop_back_val()))
+        return false;
+    }
+    return true;
+  }
+  template <typename ITy> bool match(ITy *V) { return check(V); }
+};
+
+/// Match an arbitrary undef constant. This matches poison as well.
+/// If this is an aggregate and contains a non-aggregate element that is
+/// neither undef nor poison, the aggregate is not matched.
+inline auto m_Undef() { return undef_match(); }
 
 /// Match an arbitrary poison constant.
 inline class_match<PoisonValue> m_Poison() { return class_match<PoisonValue>(); }
 
 /// Match an arbitrary Constant and ignore it.
 inline class_match<Constant> m_Constant() { return class_match<Constant>(); }
 
 /// Match an arbitrary ConstantInt and ignore it.

llvm/test/Transforms/InstSimplify/icmp-constant.ll

 bb3:
   %m = mul nuw i8 %x, 0
   br label %bb2
 }
 
 
 define <2 x i1> @heterogeneous_constvector(<2 x i8> %x) {
 ; CHECK-LABEL: @heterogeneous_constvector(
-; CHECK-NEXT:    [[C:%.*]] = icmp ult <2 x i8> [[X:%.*]], <i8 undef, i8 poison>
-; CHECK-NEXT:    ret <2 x i1> [[C]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %c = icmp ult <2 x i8> %x, <i8 undef, i8 poison>
   ret <2 x i1> %c
 }

llvm/unittests/IR/PatternMatch.cpp

EXPECT_TRUE(match(
       m_Shuffle(m_InsertElt(m_Undef(), m_SpecificInt(2), m_Zero()),
                 m_Undef(), m_ZeroMask())));
   EXPECT_TRUE(match(
       SP2, m_Shuffle(m_InsertElt(m_Undef(), m_Value(A), m_Zero()),
                      m_Undef(), m_ZeroMask())));
   EXPECT_TRUE(A == Val);
 }
 
+TEST_F(PatternMatchTest, UndefPoisonMix) {
+  Type *ScalarTy = IRB.getInt8Ty();
+  ArrayType *ArrTy = ArrayType::get(ScalarTy, 2);
+  StructType *StTy = StructType::get(ScalarTy, ScalarTy);
+  StructType *StTy2 = StructType::get(ScalarTy, StTy);
+  StructType *StTy3 = StructType::get(StTy, ScalarTy);
+  Constant *Zero = ConstantInt::getNullValue(ScalarTy);
+  UndefValue *U = UndefValue::get(ScalarTy);
+  UndefValue *P = PoisonValue::get(ScalarTy);
+
+  EXPECT_TRUE(match(ConstantVector::get({U, P}), m_Undef()));
+  EXPECT_TRUE(match(ConstantVector::get({P, U}), m_Undef()));
+
+  EXPECT_TRUE(match(ConstantArray::get(ArrTy, {U, P}), m_Undef()));
+  EXPECT_TRUE(match(ConstantArray::get(ArrTy, {P, U}), m_Undef()));
+
+  auto *UP = ConstantStruct::get(StTy, {U, P});
+  EXPECT_TRUE(match(ConstantStruct::get(StTy2, {U, UP}), m_Undef()));
+  EXPECT_TRUE(match(ConstantStruct::get(StTy2, {P, UP}), m_Undef()));
+  EXPECT_TRUE(match(ConstantStruct::get(StTy3, {UP, U}), m_Undef()));
+  EXPECT_TRUE(match(ConstantStruct::get(StTy3, {UP, P}), m_Undef()));
+
+  EXPECT_FALSE(match(ConstantStruct::get(StTy, {U, Zero}), m_Undef()));
+  EXPECT_FALSE(match(ConstantStruct::get(StTy, {Zero, U}), m_Undef()));
+  EXPECT_FALSE(match(ConstantStruct::get(StTy, {P, Zero}), m_Undef()));
+  EXPECT_FALSE(match(ConstantStruct::get(StTy, {Zero, P}), m_Undef()));
+
+  EXPECT_FALSE(match(ConstantStruct::get(StTy2, {Zero, UP}), m_Undef()));
+  EXPECT_FALSE(match(ConstantStruct::get(StTy3, {UP, Zero}), m_Undef()));
+}
+
 TEST_F(PatternMatchTest, VectorUndefInt) {
   Type *ScalarTy = IRB.getInt8Ty();
   Type *VectorTy = FixedVectorType::get(ScalarTy, 4);
   Constant *ScalarUndef = UndefValue::get(ScalarTy);
   Constant *VectorUndef = UndefValue::get(VectorTy);
   Constant *ScalarZero = Constant::getNullValue(ScalarTy);
   Constant *VectorZero = Constant::getNullValue(VectorTy);