[InstCombine][ValueTracking] Match non-uniform constant power-of-two vectors

Generalize existing constant matching to work with non-uniform constant vectors as well.

Differential Revision: https://reviews.llvm.org/D42818

llvm-svn: 324369

llvm/include/llvm/IR/PatternMatch.h

@@ -289,10 +289,27 @@ 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 (V->getType()->isVectorTy())
-      if (const auto *C = dyn_cast<Constant>(V))
+    if (V->getType()->isVectorTy()) {
+      if (const auto *C = dyn_cast<Constant>(V)) {
         if (const auto *CI = dyn_cast_or_null<ConstantInt>(C->getSplatValue()))
           return this->isValue(CI->getValue());
+
+        // Non-splat vector constant: check each element for a match.
+        unsigned NumElts = V->getType()->getVectorNumElements();
+        assert(NumElts != 0 && "Constant vector with no elements?");
+        for (unsigned i = 0; i != NumElts; ++i) {
+          Constant *Elt = C->getAggregateElement(i);
+          if (!Elt)
+            return false;
+          if (isa<UndefValue>(Elt))
+            continue;
+          auto *CI = dyn_cast<ConstantInt>(Elt);
+          if (!CI || !this->isValue(CI->getValue()))
+            return false;
+        }
+        return true;
+      }
+    }
     return false;
   }
 };
@@ -330,6 +347,14 @@ struct is_power2 {
 inline cst_pred_ty<is_power2> m_Power2() { return cst_pred_ty<is_power2>(); }
 inline api_pred_ty<is_power2> m_Power2(const APInt *&V) { return V; }
 
+struct is_power2_or_zero {
+  bool isValue(const APInt &C) { return !C || C.isPowerOf2(); }
+};
+
+/// \brief Match an integer or vector of zero or power of 2 values.
+inline cst_pred_ty<is_power2_or_zero> m_Power2OrZero() { return cst_pred_ty<is_power2_or_zero>(); }
+inline api_pred_ty<is_power2_or_zero> m_Power2OrZero(const APInt *&V) { return V; }
+
 struct is_maxsignedvalue {
   bool isValue(const APInt &C) { return C.isMaxSignedValue(); }
 };

llvm/lib/Analysis/ValueTracking.cpp

@@ -1646,14 +1646,11 @@ bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
                             const Query &Q) {
   assert(Depth <= MaxDepth && "Limit Search Depth");
 
-  if (const Constant *C = dyn_cast<Constant>(V)) {
-    if (C->isNullValue())
-      return OrZero;
-
-    const APInt *ConstIntOrConstSplatInt;
-    if (match(C, m_APInt(ConstIntOrConstSplatInt)))
-      return ConstIntOrConstSplatInt->isPowerOf2();
-  }
+  // Attempt to match against constants.
+  if (OrZero && match(V, m_Power2OrZero()))
+      return true;
+  if (match(V, m_Power2()))
+      return true;
 
   // 1 << X is clearly a power of two if the one is not shifted off the end.  If
   // it is shifted off the end then the result is undefined.

llvm/test/Transforms/InstCombine/vector-urem.ll

@@ -12,7 +12,7 @@ define <4 x i32> @test_v4i32_splatconst_pow2(<4 x i32> %a0) {
 
 define <4 x i32> @test_v4i32_const_pow2(<4 x i32> %a0) {
 ; CHECK-LABEL: @test_v4i32_const_pow2(
-; CHECK-NEXT:    [[TMP1:%.*]] = urem <4 x i32> [[A0:%.*]], <i32 1, i32 2, i32 4, i32 8>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <4 x i32> [[A0:%.*]], <i32 0, i32 1, i32 3, i32 7>
 ; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %1 = urem <4 x i32> %a0, <i32 1, i32 2, i32 4, i32 8>