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

[Valeriy Savchenko via llvm-commits]

================
@@ -3790,6 +3793,139 @@ static Instruction *foldCtpopPow2Test(ICmpInst &I, IntrinsicInst *CtpopLhs,
   return nullptr;
 }
 
+static bool isNullOrPoison(Constant *C) {
+  const auto Pred = [](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 (Pred(VecC->getOperand(I)))
+        return true;
+    }
+    return false;
+  }
+
+  if (auto *DataVec = dyn_cast<ConstantDataVector>(C)) {
+    for (unsigned I = 0; I < DataVec->getNumElements(); ++I) {
+      if (Pred(DataVec->getElementAsConstant(I)))
+        return true;
+    }
+    return false;
+  }
+
+  return Pred(C);
+}
+
+Instruction *
+InstCombinerImpl::foldICmpEqZeroVectorReduceIntrinsic(ICmpInst &Cmp,
+                                                      IntrinsicInst *II) {
+  // 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.
+  Value *InnerOp = II->getArgOperand(0);
+  const ICmpInst::Predicate Pred = Cmp.getPredicate();
+
+  // TODO: fixed vector type might be too restrictive
+  if (!II->hasOneUse() || !InnerOp->hasOneUse() ||
+      !isa<FixedVectorType>(InnerOp->getType()))
+    return nullptr;
+
+  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))) &&
+         !isNullOrPoison(C)) ||               // Case 2
+        match(InnerOp, m_ZExt(m_Value(X))) || // Case 3
+        match(InnerOp, m_SExt(m_Value(X))) || // Case 4
+        match(InnerOp, m_Neg(m_Value(X)))     // Case 5
+        ))
+    return nullptr;
+
+  SimplifyQuery S = SQ.getWithInstruction(&Cmp);
+  assert(isa<FixedVectorType>(X->getType()) && "Unexpected type");
+  auto *XTy = cast<FixedVectorType>(X->getType());
+
+  // Check for domain constraints for all supported reductions.
+  //
+  //  a. OR X_i   - has property 1  for every X
+  //  b. UMAX X_i - has property 1  for every X
+  //  c. UMIN X_i - has property 1' for every X
+  //  d. SMAX X_i - has property 1  for X >= 0
+  //  e. SMIN X_i - has property 1' for X >= 0
+  //  f. ADD X_i  - has property 1  for X >= 0 && ADD X_i doesn't sign wrap
+  //
+  // In order, for the proof to work, we need 1 (or 1') to be true for both
+  // OP f(X_i) and OP X_i and that's why below we check constraints twice.
+  //
+  // NOTE: ADD X_i holds property 1 for a mirror case as well, i.e. when
+  //       X <= 0 && ADD X_i doesn't sign wrap. However, due to the nature
+  //       of known bits, we can't reasonably hold knowledge of "either 0
+  //       or negative".
+  switch (II->getIntrinsicID()) {
+  case Intrinsic::vector_reduce_add: {
+    // If we manage to prove that the result of ADD still has at least one
+    // leading zero, we can be positive that the wrapping didn't happen for
+    // ADD f(X_i)
+    if (!isKnownNonNegative(II, S))
----------------
SavchenkoValeriy wrote:

I replaced it with direct comparison against lost bits here as well.

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