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

[RKSimon]

Summary

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

[RKSimon]

ping?

[spatel]

include/llvm/IR/PatternMatch.h

296 ↗

(On Diff #132474)

Is there a reason to limit the matcher to ConstantDataVector? If not, we should generalize to:

// 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;

The difference is that this matcher can handle weird types (and undef elts):

define <3 x i30> @test_weird_vec_type_nonsplat_pow2(<3 x i30> %a0) {
  %1 = urem <3 x i30> %a0, <i30 1, i30 2, i30 8>
  ret <3 x i30> %1
}

Side note: I copied that loop from a chunk of InstCombineCompares.cpp, and it's not the first time. We should make that available from one place, so we don't keep copying code.

[RKSimon]

Updated based on @spatel 's feedback

[spatel]

LGTM.

llvm/trunk/include/llvm/IR/PatternMatch.h

 }
 
 /// \brief This helper class is used to match scalar and vector constants that
 /// satisfy a specified predicate.
 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;
   }
 };
 
 /// \brief This helper class is used to match scalar and vector constants that
 /// satisfy a specified predicate, and bind them to an APInt.
 template <typename Predicate> struct api_pred_ty : public Predicate {
   const APInt *&Res;

 struct is_power2 {
   bool isValue(const APInt &C) { return C.isPowerOf2(); }
 };
 
 /// \brief Match an integer or vector power of 2.
 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(); }
 };
 
 inline cst_pred_ty<is_maxsignedvalue> m_MaxSignedValue() { return cst_pred_ty<is_maxsignedvalue>(); }
 inline api_pred_ty<is_maxsignedvalue> m_MaxSignedValue(const APInt *&V) { return V; }
 
 template <typename Class> struct bind_ty {

llvm/trunk/lib/Analysis/ValueTracking.cpp

 /// Return true if the given value is known to have exactly one
 /// bit set when defined. For vectors return true if every element is known to
 /// be a power of two when defined. Supports values with integer or pointer
 /// types and vectors of integers.
 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.
   if (match(V, m_Shl(m_One(), m_Value())))
     return true;
 
   // (signmask) >>l X is clearly a power of two if the one is not shifted off
   // the bottom.  If it is shifted off the bottom then the result is undefined.

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define <4 x i32> @test_v4i32_splatconst_pow2(<4 x i32> %a0) {
 ; CHECK-LABEL: @test_v4i32_splatconst_pow2(
 ; CHECK-NEXT:    [[TMP1:%.*]] = and <4 x i32> [[A0:%.*]], <i32 1, i32 1, i32 1, i32 1>
 ; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %1 = urem <4 x i32> %a0, <i32 2, i32 2, i32 2, i32 2>
   ret <4 x i32> %1
 }
 
 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>
   ret <4 x i32> %1
 }