add stripNullTest and update isKnownNonZero

Author: el-ev

llvm/include/llvm/Analysis/ValueTracking.h

@@ -999,6 +999,11 @@ LLVM_ABI void
 findValuesAffectedByCondition(Value *Cond, bool IsAssume,
                               function_ref<void(Value *)> InsertAffected);
 
+/// Returns the inner value X if the expression has the form f(X)
+/// where f(X) == 0 if and only if X == 0, otherwise returns nullptr.
+LLVM_ABI Value *stripNullTest(Value *V);
+LLVM_ABI const Value *stripNullTest(const Value *V);
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_VALUETRACKING_H

llvm/lib/Analysis/ValueTracking.cpp

@@ -3521,6 +3521,9 @@ bool isKnownNonZero(const Value *V, const APInt &DemandedElts,
       isKnownNonNullFromDominatingCondition(V, Q.CxtI, Q.DT))
     return true;
 
+  if (const Value *Stripped = stripNullTest(V))
+    return isKnownNonZero(Stripped, DemandedElts, Q, Depth);
+
   return false;
 }
 
@@ -10170,3 +10173,26 @@ void llvm::findValuesAffectedByCondition(
     }
   }
 }
+
+const Value *llvm::stripNullTest(const Value *V) {
+  // (X >> C) or/add (X & mask(C) != 0)
+  if (const auto *BO = dyn_cast<BinaryOperator>(V)) {
+    if (BO->getOpcode() == Instruction::Add ||
+        BO->getOpcode() == Instruction::Or) {
+      const Value *X;
+      const APInt *C1, *C2;
+      if (match(BO, m_c_BinOp(m_LShr(m_Value(X), m_APInt(C1)),
+                              m_ZExt(m_SpecificICmp(
+                                  ICmpInst::ICMP_NE,
+                                  m_And(m_Deferred(X), m_LowBitMask(C2)),
+                                  m_Zero())))) &&
+          C2->popcount() == C1->getZExtValue())
+        return X;
+    }
+  }
+  return nullptr;
+}
+
+Value *llvm::stripNullTest(Value *V) {
+  return const_cast<Value *>(stripNullTest(const_cast<const Value *>(V)));
+}

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

@@ -1299,22 +1299,11 @@ Instruction *InstCombinerImpl::foldICmpWithZero(ICmpInst &Cmp) {
     // will fold to a constant elsewhere.
   }
 
-  // icmp eq/ne ((X >> C) or/add (X & mask(C) != 0)), 0 -> icmp eq/ne X, 0
-  if (ICmpInst::isEquality(Pred)) {
-    auto *BO = dyn_cast<BinaryOperator>(Cmp.getOperand(0));
-    Value *X;
-    const APInt *C1, *C2;
-    if (BO &&
-        (BO->getOpcode() == Instruction::Add ||
-         BO->getOpcode() == Instruction::Or) &&
-        match(BO, m_c_BinOp(m_LShr(m_Value(X), m_APInt(C1)),
-                            m_ZExt(m_SpecificICmp(
-                                ICmpInst::ICMP_NE,
-                                m_And(m_Deferred(X), m_LowBitMask(C2)),
-                                m_Zero())))) &&
-        C2->popcount() == C1->getZExtValue())
-      return new ICmpInst(Pred, X, ConstantInt::getNullValue(X->getType()));
-  }
+  // (icmp eq/ne f(X), 0) -> (icmp eq/ne X, 0)
+  // where f(X) == 0 if and only if X == 0
+  if (ICmpInst::isEquality(Pred))
+    if (Value *Stripped = stripNullTest(Cmp.getOperand(0)))
+      return new ICmpInst(Pred, Stripped, Cmp.getOperand(1));
 
   return nullptr;
 }

llvm/test/Transforms/InstCombine/ceil-shift.ll

@@ -275,3 +275,34 @@ define i1 @ceil_shift_not_add_or(i32 %arg0) {
   %res = icmp eq i32 %quot_and_rem, 0
   ret i1 %res
 }
+
+define i32 @ceil_shift_should_infer_ge_zero(i32 %x) {
+; CHECK-LABEL: define i32 @ceil_shift_should_infer_ge_zero(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[COND_NOT:%.*]] = icmp eq i32 [[X]], 0
+; CHECK-NEXT:    br i1 [[COND_NOT]], label %[[IF_ELSE:.*]], label %[[IF_THEN:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X]], 20
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[X]], 1048575
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne i32 [[TMP2]], 0
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i1 [[TMP3]] to i32
+; CHECK-NEXT:    [[TMP5:%.*]] = add nuw nsw i32 [[TMP1]], [[TMP4]]
+; CHECK-NEXT:    ret i32 [[TMP5]]
+; CHECK:       [[IF_ELSE]]:
+; CHECK-NEXT:    ret i32 0
+;
+  %cond = icmp ne i32 %x, 0
+  br i1 %cond, label %if.then, label %if.else
+
+if.then:
+  %quot = lshr i32 %x, 20
+  %rem = and i32 %x, 1048575
+  %has_rem = icmp ne i32 %rem, 0
+  %zext_has_rem = zext i1 %has_rem to i32
+  %ceil = add nuw nsw i32 %quot, %zext_has_rem
+  %max = call i32 @llvm.umax.i32(i32 %ceil, i32 1)
+  ret i32 %max
+
+if.else:
+  ret i32 0
+}