[SVE] Make cstfp_pred_ty and cst_pred_ty work with scalable splats

[ctetreau]

add comment for mul.ll

[efriedma]

LGTM

[efriedma]

Actually, can you add similar tests that don't involve zero constants? I'm afraid at some point, we'll start folding them to zeroinitializer and end up on a different codepath.

[ctetreau]

Add unit tests

[ctetreau]

Apparently SEC is defined in sys/time.h on solaris

[efriedma]

LGTM

[ctetreau]

fix comments

[ctetreau]

rebase

llvm/include/llvm/IR/PatternMatch.h

template <int64_t Val> struct constantint_match {
   }
 };
 
 /// Match a ConstantInt with a specific value.
 template <int64_t Val> inline constantint_match<Val> m_ConstantInt() {
   return constantint_match<Val>();
 }
 
-/// This helper class is used to match scalar and fixed width vector integer
-/// constants that satisfy a specified predicate.
-/// For vector constants, undefined elements are ignored.
+/// This helper class is used to match integer constant scalars, vector splats,
+/// and fixed width vectors that satisfy a specified predicate.
+/// For fixed width vector constants, undefined elements are ignored.
 template <typename Predicate> struct cst_pred_ty : public Predicate {
   template <typename ITy> bool match(ITy *V) {
     if (const auto *CI = dyn_cast<ConstantInt>(V))
       return this->isValue(CI->getValue());
-    if (const auto *FVTy = dyn_cast<FixedVectorType>(V->getType())) {
+    if (const auto *VTy = dyn_cast<VectorType>(V->getType())) {
       if (const auto *C = dyn_cast<Constant>(V)) {
         if (const auto *CI = dyn_cast_or_null<ConstantInt>(C->getSplatValue()))
           return this->isValue(CI->getValue());
 
+        // Number of elements of a scalable vector unknown at compile time
+        auto *FVTy = dyn_cast<FixedVectorType>(VTy);
+        if (!FVTy)
+          return false;
+
         // Non-splat vector constant: check each element for a match.
         unsigned NumElts = FVTy->getNumElements();
         assert(NumElts != 0 && "Constant vector with no elements?");
         bool HasNonUndefElements = false;
         for (unsigned i = 0; i != NumElts; ++i) {
           Constant *Elt = C->getAggregateElement(i);
           if (!Elt)
             return false;
if (V->getType()->isVectorTy())
             Res = &CI->getValue();
             return true;
           }
 
     return false;
   }
 };
 
-/// This helper class is used to match scalar and vector floating-point
-/// constants that satisfy a specified predicate.
-/// For vector constants, undefined elements are ignored.
+/// This helper class is used to match float constant scalars, vector splats,
+/// and fixed width vectors that satisfy a specified predicate.
+/// For fixed width vector constants, undefined elements are ignored.
 template <typename Predicate> struct cstfp_pred_ty : public Predicate {
   template <typename ITy> bool match(ITy *V) {
     if (const auto *CF = dyn_cast<ConstantFP>(V))
       return this->isValue(CF->getValueAPF());
-    if (V->getType()->isVectorTy()) {
+    if (const auto *VTy = dyn_cast<VectorType>(V->getType())) {
       if (const auto *C = dyn_cast<Constant>(V)) {
         if (const auto *CF = dyn_cast_or_null<ConstantFP>(C->getSplatValue()))
           return this->isValue(CF->getValueAPF());
 
         // Number of elements of a scalable vector unknown at compile time
-        if (isa<ScalableVectorType>(V->getType()))
+        auto *FVTy = dyn_cast<FixedVectorType>(VTy);
+        if (!FVTy)
           return false;
 
         // Non-splat vector constant: check each element for a match.
-        unsigned NumElts = cast<VectorType>(V->getType())->getNumElements();
+        unsigned NumElts = FVTy->getNumElements();
         assert(NumElts != 0 && "Constant vector with no elements?");
         bool HasNonUndefElements = false;
         for (unsigned i = 0; i != NumElts; ++i) {
           Constant *Elt = C->getAggregateElement(i);
           if (!Elt)
             return false;
           if (isa<UndefValue>(Elt))
             continue;

llvm/test/Transforms/InstCombine/fmul.ll

 ; CHECK-NEXT:    [[MUL:%.*]] = select fast i1 [[C:%.*]], double [[X]], double [[SQR]]
 ; CHECK-NEXT:    ret double [[MUL]]
 ;
   %sqr = call double @llvm.sqrt.f64(double %x)
   %sel = select i1 %c, double %sqr, double 1.0
   %mul = fmul fast double %sqr, %sel
   ret double %mul
 }
+
+; fastmath => z * splat(0) = splat(0), even for scalable vectors
+define <vscale x 2 x float> @mul_scalable_splat_zero(<vscale x 2 x float> %z) {
+; CHECK-LABEL: @mul_scalable_splat_zero(
+; CHECK-NEXT:    ret <vscale x 2 x float> zeroinitializer
+  %shuf = shufflevector <vscale x 2 x float> insertelement (<vscale x 2 x float> undef, float 0.0, i32 0), <vscale x 2 x float> undef, <vscale x 2 x i32> zeroinitializer
+  %t3 = fmul fast <vscale x 2 x float> %shuf, %z
+  ret <vscale x 2 x float> %t3
+}

llvm/test/Transforms/InstCombine/mul.ll

 ; CHECK-NEXT:    ret <4 x i32> [[M]]
 ;
   %c = icmp slt <4 x i32> %0, zeroinitializer
   %s = sub nsw <4 x i32> zeroinitializer, %0
   %r = select <4 x i1> %c, <4 x i32> %0, <4 x i32> %s
   %m = mul <4 x i32> %r, %r
   ret <4 x i32> %m
 }
+
+; z * splat(0) = splat(0), even for scalable vectors
+define <vscale x 2 x i64> @mul_scalable_splat_zero(<vscale x 2 x i64> %z) {
+; CHECK-LABEL: @mul_scalable_splat_zero(
+; CHECK-NEXT:    ret <vscale x 2 x i64> zeroinitializer
+  %shuf = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 0, i32 0), <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
+  %t3 = mul <vscale x 2 x i64> %shuf, %z
+  ret <vscale x 2 x i64> %t3
+}

llvm/unittests/IR/PatternMatch.cpp

EXPECT_TRUE(
             m_Intrinsic<Intrinsic::instrprof_increment_step>(
                 m_Value(), m_Value(), m_Value(), m_Value(), m_SpecificInt(3))));
   EXPECT_FALSE(
       match(Intrinsic5, m_Intrinsic<Intrinsic::instrprof_increment_step>(
                             m_Value(), m_Value(), m_Value(), m_Value(),
                             m_SpecificInt(10))));
 }
 
+namespace {
+
+struct is_unsigned_zero_pred {
+  bool isValue(const APInt &C) { return C.isNullValue(); }
+};
+
+struct is_float_zero_pred {
+  bool isValue(const APFloat &C) { return C.isZero(); }
+};
+
+template <typename T> struct always_true_pred {
+  bool isValue(const T &) { return true; }
+};
+
+template <typename T> struct always_false_pred {
+  bool isValue(const T &) { return false; }
+};
+
+struct is_unsigned_max_pred {
+  bool isValue(const APInt &C) { return C.isMaxValue(); }
+};
+
+struct is_float_nan_pred {
+  bool isValue(const APFloat &C) { return C.isNaN(); }
+};
+
+} // namespace
+
+TEST_F(PatternMatchTest, ConstantPredicateType) {
+
+  // Scalar integer
+  APInt U32Max = APInt::getAllOnesValue(32);
+  APInt U32Zero = APInt::getNullValue(32);
+  APInt U32DeadBeef(32, 0xDEADBEEF);
+
+  Type *U32Ty = Type::getInt32Ty(Ctx);
+
+  Constant *CU32Max = Constant::getIntegerValue(U32Ty, U32Max);
+  Constant *CU32Zero = Constant::getIntegerValue(U32Ty, U32Zero);
+  Constant *CU32DeadBeef = Constant::getIntegerValue(U32Ty, U32DeadBeef);
+
+  EXPECT_TRUE(match(CU32Max, cst_pred_ty<is_unsigned_max_pred>()));
+  EXPECT_FALSE(match(CU32Max, cst_pred_ty<is_unsigned_zero_pred>()));
+  EXPECT_TRUE(match(CU32Max, cst_pred_ty<always_true_pred<APInt>>()));
+  EXPECT_FALSE(match(CU32Max, cst_pred_ty<always_false_pred<APInt>>()));
+
+  EXPECT_FALSE(match(CU32Zero, cst_pred_ty<is_unsigned_max_pred>()));
+  EXPECT_TRUE(match(CU32Zero, cst_pred_ty<is_unsigned_zero_pred>()));
+  EXPECT_TRUE(match(CU32Zero, cst_pred_ty<always_true_pred<APInt>>()));
+  EXPECT_FALSE(match(CU32Zero, cst_pred_ty<always_false_pred<APInt>>()));
+
+  EXPECT_FALSE(match(CU32DeadBeef, cst_pred_ty<is_unsigned_max_pred>()));
+  EXPECT_FALSE(match(CU32DeadBeef, cst_pred_ty<is_unsigned_zero_pred>()));
+  EXPECT_TRUE(match(CU32DeadBeef, cst_pred_ty<always_true_pred<APInt>>()));
+  EXPECT_FALSE(match(CU32DeadBeef, cst_pred_ty<always_false_pred<APInt>>()));
+
+  // Scalar float
+  APFloat F32NaN = APFloat::getNaN(APFloat::IEEEsingle());
+  APFloat F32Zero = APFloat::getZero(APFloat::IEEEsingle());
+  APFloat F32Pi(3.14f);
+
+  Type *F32Ty = Type::getFloatTy(Ctx);
+
+  Constant *CF32NaN = ConstantFP::get(F32Ty, F32NaN);
+  Constant *CF32Zero = ConstantFP::get(F32Ty, F32Zero);
+  Constant *CF32Pi = ConstantFP::get(F32Ty, F32Pi);
+
+  EXPECT_TRUE(match(CF32NaN, cstfp_pred_ty<is_float_nan_pred>()));
+  EXPECT_FALSE(match(CF32NaN, cstfp_pred_ty<is_float_zero_pred>()));
+  EXPECT_TRUE(match(CF32NaN, cstfp_pred_ty<always_true_pred<APFloat>>()));
+  EXPECT_FALSE(match(CF32NaN, cstfp_pred_ty<always_false_pred<APFloat>>()));
+
+  EXPECT_FALSE(match(CF32Zero, cstfp_pred_ty<is_float_nan_pred>()));
+  EXPECT_TRUE(match(CF32Zero, cstfp_pred_ty<is_float_zero_pred>()));
+  EXPECT_TRUE(match(CF32Zero, cstfp_pred_ty<always_true_pred<APFloat>>()));
+  EXPECT_FALSE(match(CF32Zero, cstfp_pred_ty<always_false_pred<APFloat>>()));
+
+  EXPECT_FALSE(match(CF32Pi, cstfp_pred_ty<is_float_nan_pred>()));
+  EXPECT_FALSE(match(CF32Pi, cstfp_pred_ty<is_float_zero_pred>()));
+  EXPECT_TRUE(match(CF32Pi, cstfp_pred_ty<always_true_pred<APFloat>>()));
+  EXPECT_FALSE(match(CF32Pi, cstfp_pred_ty<always_false_pred<APFloat>>()));
+
+  ElementCount FixedEC(4, false);
+  ElementCount ScalableEC(4, true);
+
+  // Vector splat
+
+  for (auto EC : {FixedEC, ScalableEC}) {
+    // integer
+
+    Constant *CSplatU32Max = ConstantVector::getSplat(EC, CU32Max);
+    Constant *CSplatU32Zero = ConstantVector::getSplat(EC, CU32Zero);
+    Constant *CSplatU32DeadBeef = ConstantVector::getSplat(EC, CU32DeadBeef);
+
+    EXPECT_TRUE(match(CSplatU32Max, cst_pred_ty<is_unsigned_max_pred>()));
+    EXPECT_FALSE(match(CSplatU32Max, cst_pred_ty<is_unsigned_zero_pred>()));
+    EXPECT_TRUE(match(CSplatU32Max, cst_pred_ty<always_true_pred<APInt>>()));
+    EXPECT_FALSE(match(CSplatU32Max, cst_pred_ty<always_false_pred<APInt>>()));
+
+    EXPECT_FALSE(match(CSplatU32Zero, cst_pred_ty<is_unsigned_max_pred>()));
+    EXPECT_TRUE(match(CSplatU32Zero, cst_pred_ty<is_unsigned_zero_pred>()));
+    EXPECT_TRUE(match(CSplatU32Zero, cst_pred_ty<always_true_pred<APInt>>()));
+    EXPECT_FALSE(match(CSplatU32Zero, cst_pred_ty<always_false_pred<APInt>>()));
+
+    EXPECT_FALSE(match(CSplatU32DeadBeef, cst_pred_ty<is_unsigned_max_pred>()));
+    EXPECT_FALSE(
+        match(CSplatU32DeadBeef, cst_pred_ty<is_unsigned_zero_pred>()));
+    EXPECT_TRUE(
+        match(CSplatU32DeadBeef, cst_pred_ty<always_true_pred<APInt>>()));
+    EXPECT_FALSE(
+        match(CSplatU32DeadBeef, cst_pred_ty<always_false_pred<APInt>>()));
+
+    // float
+
+    Constant *CSplatF32NaN = ConstantVector::getSplat(EC, CF32NaN);
+    Constant *CSplatF32Zero = ConstantVector::getSplat(EC, CF32Zero);
+    Constant *CSplatF32Pi = ConstantVector::getSplat(EC, CF32Pi);
+
+    EXPECT_TRUE(match(CSplatF32NaN, cstfp_pred_ty<is_float_nan_pred>()));
+    EXPECT_FALSE(match(CSplatF32NaN, cstfp_pred_ty<is_float_zero_pred>()));
+    EXPECT_TRUE(
+        match(CSplatF32NaN, cstfp_pred_ty<always_true_pred<APFloat>>()));
+    EXPECT_FALSE(
+        match(CSplatF32NaN, cstfp_pred_ty<always_false_pred<APFloat>>()));
+
+    EXPECT_FALSE(match(CSplatF32Zero, cstfp_pred_ty<is_float_nan_pred>()));
+    EXPECT_TRUE(match(CSplatF32Zero, cstfp_pred_ty<is_float_zero_pred>()));
+    EXPECT_TRUE(
+        match(CSplatF32Zero, cstfp_pred_ty<always_true_pred<APFloat>>()));
+    EXPECT_FALSE(
+        match(CSplatF32Zero, cstfp_pred_ty<always_false_pred<APFloat>>()));
+
+    EXPECT_FALSE(match(CSplatF32Pi, cstfp_pred_ty<is_float_nan_pred>()));
+    EXPECT_FALSE(match(CSplatF32Pi, cstfp_pred_ty<is_float_zero_pred>()));
+    EXPECT_TRUE(match(CSplatF32Pi, cstfp_pred_ty<always_true_pred<APFloat>>()));
+    EXPECT_FALSE(
+        match(CSplatF32Pi, cstfp_pred_ty<always_false_pred<APFloat>>()));
+  }
+
+  // Int arbitrary vector
+
+  Constant *CMixedU32 = ConstantVector::get({CU32Max, CU32Zero, CU32DeadBeef});
+  Constant *CU32Undef = UndefValue::get(U32Ty);
+  Constant *CU32MaxWithUndef =
+      ConstantVector::get({CU32Undef, CU32Max, CU32Undef});
+
+  EXPECT_FALSE(match(CMixedU32, cst_pred_ty<is_unsigned_max_pred>()));
+  EXPECT_FALSE(match(CMixedU32, cst_pred_ty<is_unsigned_zero_pred>()));
+  EXPECT_TRUE(match(CMixedU32, cst_pred_ty<always_true_pred<APInt>>()));
+  EXPECT_FALSE(match(CMixedU32, cst_pred_ty<always_false_pred<APInt>>()));
+
+  EXPECT_TRUE(match(CU32MaxWithUndef, cst_pred_ty<is_unsigned_max_pred>()));
+  EXPECT_FALSE(match(CU32MaxWithUndef, cst_pred_ty<is_unsigned_zero_pred>()));
+  EXPECT_TRUE(match(CU32MaxWithUndef, cst_pred_ty<always_true_pred<APInt>>()));
+  EXPECT_FALSE(
+      match(CU32MaxWithUndef, cst_pred_ty<always_false_pred<APInt>>()));
+
+  // Float arbitrary vector
+
+  Constant *CMixedF32 = ConstantVector::get({CF32NaN, CF32Zero, CF32Pi});
+  Constant *CF32Undef = UndefValue::get(F32Ty);
+  Constant *CF32NaNWithUndef =
+      ConstantVector::get({CF32Undef, CF32NaN, CF32Undef});
+
+  EXPECT_FALSE(match(CMixedF32, cstfp_pred_ty<is_float_nan_pred>()));
+  EXPECT_FALSE(match(CMixedF32, cstfp_pred_ty<is_float_zero_pred>()));
+  EXPECT_TRUE(match(CMixedF32, cstfp_pred_ty<always_true_pred<APFloat>>()));
+  EXPECT_FALSE(match(CMixedF32, cstfp_pred_ty<always_false_pred<APFloat>>()));
+
+  EXPECT_TRUE(match(CF32NaNWithUndef, cstfp_pred_ty<is_float_nan_pred>()));
+  EXPECT_FALSE(match(CF32NaNWithUndef, cstfp_pred_ty<is_float_zero_pred>()));
+  EXPECT_TRUE(
+      match(CF32NaNWithUndef, cstfp_pred_ty<always_true_pred<APFloat>>()));
+  EXPECT_FALSE(
+      match(CF32NaNWithUndef, cstfp_pred_ty<always_false_pred<APFloat>>()));
+}
+
 template <typename T> struct MutableConstTest : PatternMatchTest { };
 
 typedef ::testing::Types<std::tuple<Value*, Instruction*>,
                          std::tuple<const Value*, const Instruction *>>
     MutableConstTestTypes;
 TYPED_TEST_CASE(MutableConstTest, MutableConstTestTypes);
 
 TYPED_TEST(MutableConstTest, ICmp) {