[PatternMatch] Add m_APInt/m_APFloat matchers accepting undef

[nikic]

Summary

The current m_APInt() and m_APFloat() matchers do not accept splats that include undefs (unlike m_Zero() and other matchers for specific values). We can't simply change the default behavior, as there are existing transforms that would not be safe with undefs.

For this reason, I'm introducing new m_APIntAllowUndef() and m_APFloatAllowUndef() matchers, that allow splats with undefs. Additionally, m_APIntForbidUndef() and m_APFloatForbidUndef() are added. These have the same behavior as the existing m_APInt() and m_APFloat(), but serve as an explicit indication that undefs were explicitly considered and found unsound for this transform. This helps distinguish them from existing uses of m_APInt() where we do not know whether undefs can or cannot be allowed without additional review.

[spatel]

LGTM - @lebedev.ri may have some notes about existing code where we can use 'allow' or 'forbid' flavors explicitly. (There were earlier experiments with allowing undefs wholesale in these matchers.)

llvm/include/llvm/IR/PatternMatch.h

 /// Combine two pattern matchers matching L && R
 template <typename LTy, typename RTy>
 inline match_combine_and<LTy, RTy> m_CombineAnd(const LTy &L, const RTy &R) {
   return match_combine_and<LTy, RTy>(L, R);
 }
 
 struct apint_match {
   const APInt *&Res;
+  bool AllowUndef;
 
-  apint_match(const APInt *&R) : Res(R) {}
+  apint_match(const APInt *&Res, bool AllowUndef)
+    : Res(Res), AllowUndef(AllowUndef) {}
 
   template <typename ITy> bool match(ITy *V) {
     if (auto *CI = dyn_cast<ConstantInt>(V)) {
       Res = &CI->getValue();
       return true;
     }
     if (V->getType()->isVectorTy())
       if (const auto *C = dyn_cast<Constant>(V))
-        if (auto *CI = dyn_cast_or_null<ConstantInt>(C->getSplatValue())) {
+        if (auto *CI = dyn_cast_or_null<ConstantInt>(
+                C->getSplatValue(AllowUndef))) {
           Res = &CI->getValue();
           return true;
         }
     return false;
   }
 };
 // Either constexpr if or renaming ConstantFP::getValueAPF to
 // ConstantFP::getValue is needed to do it via single template
 // function for both apint/apfloat.
 struct apfloat_match {
   const APFloat *&Res;
-  apfloat_match(const APFloat *&R) : Res(R) {}
+  bool AllowUndef;
+
+  apfloat_match(const APFloat *&Res, bool AllowUndef)
+      : Res(Res), AllowUndef(AllowUndef) {}
+
   template <typename ITy> bool match(ITy *V) {
     if (auto *CI = dyn_cast<ConstantFP>(V)) {
       Res = &CI->getValueAPF();
       return true;
     }
     if (V->getType()->isVectorTy())
       if (const auto *C = dyn_cast<Constant>(V))
-        if (auto *CI = dyn_cast_or_null<ConstantFP>(C->getSplatValue())) {
+        if (auto *CI = dyn_cast_or_null<ConstantFP>(
+                C->getSplatValue(AllowUndef))) {
           Res = &CI->getValueAPF();
           return true;
         }
     return false;
   }
 };
 
 /// Match a ConstantInt or splatted ConstantVector, binding the
 /// specified pointer to the contained APInt.
-inline apint_match m_APInt(const APInt *&Res) { return Res; }
+inline apint_match m_APInt(const APInt *&Res) {
+  // Forbid undefs by default to maintain previous behavior.
+  return apint_match(Res, /* AllowUndef */ false);
+}
+
+/// Match APInt while allowing undefs in splat vector constants.
+inline apint_match m_APIntAllowUndef(const APInt *&Res) {
+  return apint_match(Res, /* AllowUndef */ true);
+}
+
+/// Match APInt while forbidding undefs in splat vector constants.
+inline apint_match m_APIntForbidUndef(const APInt *&Res) {
+  return apint_match(Res, /* AllowUndef */ false);
+}
 
 /// Match a ConstantFP or splatted ConstantVector, binding the
 /// specified pointer to the contained APFloat.
-inline apfloat_match m_APFloat(const APFloat *&Res) { return Res; }
+inline apfloat_match m_APFloat(const APFloat *&Res) {
+  // Forbid undefs by default to maintain previous behavior.
+  return apfloat_match(Res, /* AllowUndef */ false);
+}
+
+/// Match APFloat while allowing undefs in splat vector constants.
+inline apfloat_match m_APFloatAllowUndef(const APFloat *&Res) {
+  return apfloat_match(Res, /* AllowUndef */ true);
+}
+
+/// Match APFloat while forbidding undefs in splat vector constants.
+inline apfloat_match m_APFloatForbidUndef(const APFloat *&Res) {
+  return apfloat_match(Res, /* AllowUndef */ false);
+}
 
 template <int64_t Val> struct constantint_match {
   template <typename ITy> bool match(ITy *V) {
     if (const auto *CI = dyn_cast<ConstantInt>(V)) {
       const APInt &CIV = CI->getValue();
       if (Val >= 0)
         return CIV == static_cast<uint64_t>(Val);
       // If Val is negative, and CI is shorter than it, truncate to the right

llvm/unittests/IR/PatternMatch.cpp

TEST_F(PatternMatchTest, VectorUndefInt) {
   EXPECT_FALSE(match(VectorZero, m_Undef()));
   EXPECT_FALSE(match(VectorZeroUndef, m_Undef()));
 
   EXPECT_FALSE(match(ScalarUndef, m_Zero()));
   EXPECT_FALSE(match(VectorUndef, m_Zero()));
   EXPECT_TRUE(match(ScalarZero, m_Zero()));
   EXPECT_TRUE(match(VectorZero, m_Zero()));
   EXPECT_TRUE(match(VectorZeroUndef, m_Zero()));
+
+  const APInt *C;
+  // Regardless of whether undefs are allowed,
+  // a fully undef constant does not match.
+  EXPECT_FALSE(match(ScalarUndef, m_APInt(C)));
+  EXPECT_FALSE(match(ScalarUndef, m_APIntForbidUndef(C)));
+  EXPECT_FALSE(match(ScalarUndef, m_APIntAllowUndef(C)));
+  EXPECT_FALSE(match(VectorUndef, m_APInt(C)));
+  EXPECT_FALSE(match(VectorUndef, m_APIntForbidUndef(C)));
+  EXPECT_FALSE(match(VectorUndef, m_APIntAllowUndef(C)));
+
+  // We can always match simple constants and simple splats.
+  C = nullptr;
+  EXPECT_TRUE(match(ScalarZero, m_APInt(C)));
+  EXPECT_TRUE(C->isNullValue());
+  C = nullptr;
+  EXPECT_TRUE(match(ScalarZero, m_APIntForbidUndef(C)));
+  EXPECT_TRUE(C->isNullValue());
+  C = nullptr;
+  EXPECT_TRUE(match(ScalarZero, m_APIntAllowUndef(C)));
+  EXPECT_TRUE(C->isNullValue());
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZero, m_APInt(C)));
+  EXPECT_TRUE(C->isNullValue());
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZero, m_APIntForbidUndef(C)));
+  EXPECT_TRUE(C->isNullValue());
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZero, m_APIntAllowUndef(C)));
+  EXPECT_TRUE(C->isNullValue());
+
+  // Whether splats with undef can be matched depends on the matcher.
+  EXPECT_FALSE(match(VectorZeroUndef, m_APInt(C)));
+  EXPECT_FALSE(match(VectorZeroUndef, m_APIntForbidUndef(C)));
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZeroUndef, m_APIntAllowUndef(C)));
+  EXPECT_TRUE(C->isNullValue());
 }
 
 TEST_F(PatternMatchTest, VectorUndefFloat) {
   Type *ScalarTy = IRB.getFloatTy();
   Type *VectorTy = VectorType::get(ScalarTy, 4);
   Constant *ScalarUndef = UndefValue::get(ScalarTy);
   Constant *VectorUndef = UndefValue::get(VectorTy);
   Constant *ScalarZero = Constant::getNullValue(ScalarTy);
TEST_F(PatternMatchTest, VectorUndefFloat) {
   EXPECT_FALSE(match(VectorZero, m_Undef()));
   EXPECT_FALSE(match(VectorZeroUndef, m_Undef()));
 
   EXPECT_FALSE(match(ScalarUndef, m_AnyZeroFP()));
   EXPECT_FALSE(match(VectorUndef, m_AnyZeroFP()));
   EXPECT_TRUE(match(ScalarZero, m_AnyZeroFP()));
   EXPECT_TRUE(match(VectorZero, m_AnyZeroFP()));
   EXPECT_TRUE(match(VectorZeroUndef, m_AnyZeroFP()));
+
+  const APFloat *C;
+  // Regardless of whether undefs are allowed,
+  // a fully undef constant does not match.
+  EXPECT_FALSE(match(ScalarUndef, m_APFloat(C)));
+  EXPECT_FALSE(match(ScalarUndef, m_APFloatForbidUndef(C)));
+  EXPECT_FALSE(match(ScalarUndef, m_APFloatAllowUndef(C)));
+  EXPECT_FALSE(match(VectorUndef, m_APFloat(C)));
+  EXPECT_FALSE(match(VectorUndef, m_APFloatForbidUndef(C)));
+  EXPECT_FALSE(match(VectorUndef, m_APFloatAllowUndef(C)));
+
+  // We can always match simple constants and simple splats.
+  C = nullptr;
+  EXPECT_TRUE(match(ScalarZero, m_APFloat(C)));
+  EXPECT_TRUE(C->isZero());
+  C = nullptr;
+  EXPECT_TRUE(match(ScalarZero, m_APFloatForbidUndef(C)));
+  EXPECT_TRUE(C->isZero());
+  C = nullptr;
+  EXPECT_TRUE(match(ScalarZero, m_APFloatAllowUndef(C)));
+  EXPECT_TRUE(C->isZero());
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZero, m_APFloat(C)));
+  EXPECT_TRUE(C->isZero());
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZero, m_APFloatForbidUndef(C)));
+  EXPECT_TRUE(C->isZero());
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZero, m_APFloatAllowUndef(C)));
+  EXPECT_TRUE(C->isZero());
+
+  // Whether splats with undef can be matched depends on the matcher.
+  EXPECT_FALSE(match(VectorZeroUndef, m_APFloat(C)));
+  EXPECT_FALSE(match(VectorZeroUndef, m_APFloatForbidUndef(C)));
+  C = nullptr;
+  EXPECT_TRUE(match(VectorZeroUndef, m_APFloatAllowUndef(C)));
+  EXPECT_TRUE(C->isZero());
 }
 
 TEST_F(PatternMatchTest, FloatingPointFNeg) {
   Type *FltTy = IRB.getFloatTy();
   Value *One = ConstantFP::get(FltTy, 1.0);
   Value *Z = ConstantFP::get(FltTy, 0.0);
   Value *NZ = ConstantFP::get(FltTy, -0.0);
   Value *V = IRB.CreateFNeg(One);