[ValueTracking, InstCombine] extend isKnownToBeAPowerOfTwo() to handle vector splat constants

[spatel]

Summary

We could try harder and handle non-splat vector constants too, but that seems much rarer to me.
Note that the div test isn't resolved because there's a check for isIntegerTy() guarding that transform.

[majnemer]

lib/Analysis/ValueTracking.cpp

1503–1507

I think you could use m_APInt to merge these cases together.

[spatel]

Patch updated:
Today, I learned that m_APInt works for splats.
Thanks, David!

[majnemer]

LGTM

lib/Analysis/ValueTracking.cpp

1504

I don't think you need to initialize this.

lib/Analysis/ValueTracking.cpp

 /// 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(Value *V, bool OrZero, unsigned Depth,
                             const Query &Q) {
   if (Constant *C = dyn_cast<Constant>(V)) {
     if (C->isNullValue())
       return OrZero;
-    if (ConstantInt *CI = dyn_cast<ConstantInt>(C))
-      return CI->getValue().isPowerOf2();
-    // TODO: Handle vector constants.
+
+    const APInt *ConstIntOrConstSplatInt = nullptr;
+    if (match(C, m_APInt(ConstIntOrConstSplatInt)))
+      return ConstIntOrConstSplatInt->isPowerOf2();
   }
 
   // 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;
 
   // (signbit) >>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.

test/Transforms/InstCombine/div.ll

;
   ret i32 %mul
 }
 
 ; FIXME: Vector should get same transform as scalar.
 
 define <2 x i32> @test36vec(<2 x i32> %A) {
 ; CHECK-LABEL: @test36vec(
 ; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> %A, <i32 2147483647, i32 2147483647>
-; CHECK-NEXT:    [[SHL:%.*]] = shl nsw <2 x i32> <i32 1, i32 1>, %A
+; CHECK-NEXT:    [[SHL:%.*]] = shl nuw nsw <2 x i32> <i32 1, i32 1>, %A
 ; CHECK-NEXT:    [[MUL:%.*]] = sdiv exact <2 x i32> [[AND]], [[SHL]]
 ; CHECK-NEXT:    ret <2 x i32> [[MUL]]
 ;
   %and = and <2 x i32> %A, <i32 2147483647, i32 2147483647>
   %shl = shl nsw <2 x i32> <i32 1, i32 1>, %A
   %mul = sdiv exact <2 x i32> %and, %shl
   ret <2 x i32> %mul
 }

test/Transforms/InstCombine/rem.ll

 define i32 @test3(i32 %A) {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT: [[AND:%.*]] = and i32 %A, 7
 ; CHECK-NEXT: ret i32 [[AND]]
   %B = urem i32 %A, 8
   ret i32 %B
 }
 
-; FIXME: This could be an 'and' just like above.
-
 define <2 x i32> @vec_power_of_2_constant_splat_divisor(<2 x i32> %A) {
 ; CHECK-LABEL: @vec_power_of_2_constant_splat_divisor(
-; CHECK-NEXT:    [[B:%.*]] = urem <2 x i32> %A, <i32 8, i32 8>
+; CHECK-NEXT:    [[B:%.*]] = and <2 x i32> %A, <i32 7, i32 7>
 ; CHECK-NEXT:    ret <2 x i32> [[B]]
 ;
   %B = urem <2 x i32> %A, <i32 8, i32 8>
   ret <2 x i32> %B
 }
 
-; FIXME: And it shouldn't matter whether we have ConstantVector or ConstantDataVector.
-
 define <2 x i19> @weird_vec_power_of_2_constant_splat_divisor(<2 x i19> %A) {
 ; CHECK-LABEL: @weird_vec_power_of_2_constant_splat_divisor(
-; CHECK-NEXT:    [[B:%.*]] = urem <2 x i19> %A, <i19 8, i19 8>
+; CHECK-NEXT:    [[B:%.*]] = and <2 x i19> %A, <i19 7, i19 7>
 ; CHECK-NEXT:    ret <2 x i19> [[B]]
 ;
   %B = urem <2 x i19> %A, <i19 8, i19 8>
   ret <2 x i19> %B
 }
 
 define i1 @test3a(i32 %A) {
 ; CHECK-LABEL: @test3a(

test/Transforms/InstCombine/shift.ll

 ; CHECK-NEXT:    [[DIV2:%.*]] = udiv i32 %a, [[DIV]]
 ; CHECK-NEXT:    ret i32 [[DIV2]]
 ;
   %div = lshr i32 4096, %b    ; must be exact otherwise we'd divide by zero
   %div2 = udiv i32 %a, %div
   ret i32 %div2
 }
 
-; FIXME: Vector lshr should be treated the same as scalar. 
-
 define <2 x i32> @test42vec(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @test42vec(
-; CHECK-NEXT:    [[DIV:%.*]] = lshr <2 x i32> <i32 4096, i32 4096>, %b
+; CHECK-NEXT:    [[DIV:%.*]] = lshr exact <2 x i32> <i32 4096, i32 4096>, %b
 ; CHECK-NEXT:    [[DIV2:%.*]] = udiv <2 x i32> %a, [[DIV]]
 ; CHECK-NEXT:    ret <2 x i32> [[DIV2]]
 ;
   %div = lshr <2 x i32> <i32 4096, i32 4096>, %b    ; must be exact otherwise we'd divide by zero
   %div2 = udiv <2 x i32> %a, %div
   ret <2 x i32> %div2
 }