[llvm] [VectorCombine] Fold vector.reduce.OP(F(X)) == 0 -> OP(X) == 0 (PR #173069)

[Valeriy Savchenko via llvm-commits]

================
@@ -3806,6 +3807,197 @@ bool VectorCombine::foldCastFromReductions(Instruction &I) {
   return true;
 }
 
+/// Check whether the constant contains a null or poison element.
+static bool containsNullOrPoison(Constant *C) {
+  const auto IsNullOrPoison = [](Constant *X) {
+    return X->isNullValue() || isa<UndefValue>(X) || isa<PoisonValue>(X);
+  };
+
+  if (auto *VecC = dyn_cast<ConstantVector>(C)) {
+    for (unsigned I = 0; I < VecC->getNumOperands(); ++I) {
+      if (IsNullOrPoison(VecC->getOperand(I)))
+        return true;
+    }
+    return false;
+  }
+
+  if (auto *DataVec = dyn_cast<ConstantDataVector>(C)) {
+    for (unsigned I = 0; I < DataVec->getNumElements(); ++I) {
+      if (IsNullOrPoison(DataVec->getElementAsConstant(I)))
+        return true;
+    }
+    return false;
+  }
+
+  return IsNullOrPoison(C);
+}
+
+bool VectorCombine::foldICmpEqZeroVectorReduce(Instruction &I) {
+  // vector.reduce.OP f(X_i) == 0 -> vector.reduce.OP X_i == 0
+  //
+  // We can prove it for cases when:
+  //
+  //   1.  OP X_i == 0 <=> \forall i \in [1, N] X_i == 0
+  //   1'. OP X_i == 0 <=> \exists j \in [1, N] X_j == 0
+  //   2.  f(x) == 0 <=> x == 0
+  //
+  // From 1 and 2 (or 1' and 2), we can infer that
+  //
+  //   OP f(X_i) == 0 <=> OP X_i == 0.
+  //
+  // For some of the OP's and f's, we need to have domain constraints on X
+  // to ensure properties 1 (or 1') and 2.
+  CmpPredicate Pred;
+  Value *Op;
+  if (!match(&I, m_ICmp(Pred, m_Value(Op), m_Zero())) ||
+      !ICmpInst::isEquality(Pred))
+    return false;
+
+  auto *II = dyn_cast<IntrinsicInst>(Op);
+  if (!II)
+    return false;
+
+  switch (II->getIntrinsicID()) {
+  case Intrinsic::vector_reduce_add:
+  case Intrinsic::vector_reduce_or:
+  case Intrinsic::vector_reduce_umin:
+  case Intrinsic::vector_reduce_umax:
+  case Intrinsic::vector_reduce_smin:
+  case Intrinsic::vector_reduce_smax:
+    break;
+  default:
+    return false;
+  }
+
+  Value *InnerOp = II->getArgOperand(0);
+
+  // TODO: fixed vector type might be too restrictive
+  if (!II->hasOneUse() || !InnerOp->hasOneUse() ||
+      !isa<FixedVectorType>(InnerOp->getType()))
+    return false;
+
+  Value *X = nullptr;
+  Constant *C = nullptr;
+
+  // Check for zero-preserving operations where f(x) = 0 <=> x = 0
+  //
+  //   1. f(x) = shl nuw x, y for arbitrary y
+  //   2. f(x) = mul nuw x, c for defined c != 0
+  //   3. f(x) = zext x
+  //   4. f(x) = sext x
+  //   5. f(x) = neg x
+  //
+  if (!(match(InnerOp, m_NUWShl(m_Value(X),
+                                m_Value())) || // Case 1
+        (match(InnerOp, m_NUWMul(m_Value(X), m_Constant(C))) &&
+         !containsNullOrPoison(C)) ||         // Case 2
----------------
SavchenkoValeriy wrote:

It looks like `m_CheckedInt` accepts a predicate for `APInt` not `Constant` itself.

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