[InstCombine] canonicalize funnel shift constant shift amount to be modulo bitwidth

[spatel]

Summary

The shift argument is defined to be modulo the bitwidth, so if that argument is a constant, we can always reduce the constant to its minimal form to allow better CSE and other follow-on transforms.

There was already a test that shows a follow-on simplification, so I assume we just had not made our way around to adding this transform yet.

[lebedev.ri]

Hmm, this looks good, but i have a nit.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2000 ↗

(On Diff #190672)

Will m_Constant() also math constant expressions?
Maybe add some test that we don't do anything stupid with them?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2000 ↗

(On Diff #190672)

Yes, that'll match a ConstantExpr...and yes, we often get into trouble with those, so good catch!
I'll add some tests.

[nikic]

You might also want to kill this TODO: https://github.com/llvm-mirror/llvm/blob/master/lib/Analysis/InstructionSimplify.cpp#L4920

Technically it's not quite solved because we're not guaranteed canonical IR, but it's probably not worth dealing with this in InstSimplify if we have this in InstCombine.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2012 ↗

(On Diff #190672)

Can also drop this BitWidth.

[spatel]

In D59374#1429674, @nikic wrote:

You might also want to kill this TODO: https://github.com/llvm-mirror/llvm/blob/master/lib/Analysis/InstructionSimplify.cpp#L4920

Technically it's not quite solved because we're not guaranteed canonical IR, but it's probably not worth dealing with this in InstSimplify if we have this in InstCombine.

Agreed - there's almost no chance that other passes will run before instcombine and/or create non-canonical funnel shifts. Thanks for the quick reviews!

llvm/trunk/include/llvm/IR/Constant.h

public:
   /// Return true if this is a floating-point NaN constant or a vector
   /// floating-point constant with all NaN elements.
   bool isNaN() const;
 
   /// Return true if this is a vector constant that includes any undefined
   /// elements.
   bool containsUndefElement() const;
 
+  /// Return true if this is a vector constant that includes any constant
+  /// expressions.
+  bool containsConstantExpression() const;
+
   /// Return true if evaluation of this constant could trap. This is true for
   /// things like constant expressions that could divide by zero.
   bool canTrap() const;
 
   /// Return true if the value can vary between threads.
   bool isThreadDependent() const;
 
   /// Return true if the value is dependent on a dllimport variable.

llvm/trunk/lib/Analysis/InstructionSimplify.cpp

case Intrinsic::fshr: {
 
     // If shift amount is undef, assume it is zero.
     if (match(ShAmtArg, m_Undef()))
       return ArgBegin[IID == Intrinsic::fshl ? 0 : 1];
 
     const APInt *ShAmtC;
     if (match(ShAmtArg, m_APInt(ShAmtC))) {
       // If there's effectively no shift, return the 1st arg or 2nd arg.
-      // TODO: For vectors, we could check each element of a non-splat constant.
       APInt BitWidth = APInt(ShAmtC->getBitWidth(), ShAmtC->getBitWidth());
       if (ShAmtC->urem(BitWidth).isNullValue())
         return ArgBegin[IID == Intrinsic::fshl ? 0 : 1];
     }
     return nullptr;
   }
   default:
     return nullptr;

llvm/trunk/lib/IR/Constants.cpp

if (!getType()->isVectorTy())
     return false;
   for (unsigned i = 0, e = getType()->getVectorNumElements(); i != e; ++i)
     if (isa<UndefValue>(getAggregateElement(i)))
       return true;
 
   return false;
 }
 
+bool Constant::containsConstantExpression() const {
+  if (!getType()->isVectorTy())
+    return false;
+  for (unsigned i = 0, e = getType()->getVectorNumElements(); i != e; ++i)
+    if (isa<ConstantExpr>(getAggregateElement(i)))
+      return true;
+
+  return false;
+}
+
 /// Constructor to create a '0' constant of arbitrary type.
 Constant *Constant::getNullValue(Type *Ty) {
   switch (Ty->getTypeID()) {
   case Type::IntegerTyID:
     return ConstantInt::get(Ty, 0);
   case Type::HalfTyID:
     return ConstantFP::get(Ty->getContext(),
                            APFloat::getZero(APFloat::IEEEhalf()));

llvm/trunk/lib/Transforms/InstCombine/InstCombineCalls.cpp

Instruction *InstCombiner::visitCallInst(CallInst &CI) {
 
   case Intrinsic::ctpop:
     if (auto *I = foldCtpop(*II, *this))
       return I;
     break;
 
   case Intrinsic::fshl:
   case Intrinsic::fshr: {
+    // Canonicalize a shift amount constant operand to be modulo the bit-width.
+    unsigned BitWidth = II->getType()->getScalarSizeInBits();
+    Constant *ShAmtC;
+    if (match(II->getArgOperand(2), m_Constant(ShAmtC)) &&
+        !isa<ConstantExpr>(ShAmtC) && !ShAmtC->containsConstantExpression()) {
+      Constant *WidthC = ConstantInt::get(II->getType(), BitWidth);
+      Constant *ModuloC = ConstantExpr::getURem(ShAmtC, WidthC);
+      if (ModuloC != ShAmtC) {
+        II->setArgOperand(2, ModuloC);
+        return II;
+      }
+    }
+
     const APInt *SA;
     if (match(II->getArgOperand(2), m_APInt(SA))) {
       Value *Op0 = II->getArgOperand(0), *Op1 = II->getArgOperand(1);
-      unsigned BitWidth = SA->getBitWidth();
       uint64_t ShiftAmt = SA->urem(BitWidth);
       assert(ShiftAmt != 0 && "SimplifyCall should have handled zero shift");
       // Normalize to funnel shift left.
       if (II->getIntrinsicID() == Intrinsic::fshr)
         ShiftAmt = BitWidth - ShiftAmt;
 
       // fshl(X, 0, C) -> shl X, C
       // fshl(X, undef, C) -> shl X, C
       if (match(Op1, m_Zero()) || match(Op1, m_Undef()))
         return BinaryOperator::CreateShl(
             Op0, ConstantInt::get(II->getType(), ShiftAmt));
 
       // fshl(0, X, C) -> lshr X, (BW-C)
       // fshl(undef, X, C) -> lshr X, (BW-C)
       if (match(Op0, m_Zero()) || match(Op0, m_Undef()))
         return BinaryOperator::CreateLShr(
             Op1, ConstantInt::get(II->getType(), BitWidth - ShiftAmt));
     }
 
     // The shift amount (operand 2) of a funnel shift is modulo the bitwidth,
     // so only the low bits of the shift amount are demanded if the bitwidth is
     // a power-of-2.
-    unsigned BitWidth = II->getType()->getScalarSizeInBits();
     if (!isPowerOf2_32(BitWidth))
       break;
     APInt Op2Demanded = APInt::getLowBitsSet(BitWidth, Log2_32_Ceil(BitWidth));
     KnownBits Op2Known(BitWidth);
     if (SimplifyDemandedBits(II, 2, Op2Demanded, Op2Known))
       return &CI;
     break;
   }

llvm/trunk/test/Transforms/InstCombine/fsh.ll

 ;
   %z = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> <i31 7, i31 7>)
   %r = and <2 x i31> %z, <i31 63, i31 31>
   ret <2 x i31> %r
 }
 
 define i32 @fshl_constant_shift_amount_modulo_bitwidth(i32 %x, i32 %y) {
 ; CHECK-LABEL: @fshl_constant_shift_amount_modulo_bitwidth(
-; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.fshl.i32(i32 [[X:%.*]], i32 [[Y:%.*]], i32 33)
+; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.fshl.i32(i32 [[X:%.*]], i32 [[Y:%.*]], i32 1)
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %r = call i32 @llvm.fshl.i32(i32 %x, i32 %y, i32 33)
   ret i32 %r
 }
 
 define i33 @fshr_constant_shift_amount_modulo_bitwidth(i33 %x, i33 %y) {
 ; CHECK-LABEL: @fshr_constant_shift_amount_modulo_bitwidth(
-; CHECK-NEXT:    [[R:%.*]] = call i33 @llvm.fshr.i33(i33 [[X:%.*]], i33 [[Y:%.*]], i33 34)
+; CHECK-NEXT:    [[R:%.*]] = call i33 @llvm.fshr.i33(i33 [[X:%.*]], i33 [[Y:%.*]], i33 1)
 ; CHECK-NEXT:    ret i33 [[R]]
 ;
   %r = call i33 @llvm.fshr.i33(i33 %x, i33 %y, i33 34)
   ret i33 %r
 }
 
+@external_global = external global i8
+
+define i33 @fshr_constant_shift_amount_modulo_bitwidth_constexpr(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_constant_shift_amount_modulo_bitwidth_constexpr(
+; CHECK-NEXT:    [[R:%.*]] = call i33 @llvm.fshr.i33(i33 [[X:%.*]], i33 [[Y:%.*]], i33 ptrtoint (i8* @external_global to i33))
+; CHECK-NEXT:    ret i33 [[R]]
+;
+  %shamt = ptrtoint i8* @external_global to i33
+  %r = call i33 @llvm.fshr.i33(i33 %x, i33 %y, i33 %shamt)
+  ret i33 %r
+}
+
 define <2 x i32> @fshr_constant_shift_amount_modulo_bitwidth_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @fshr_constant_shift_amount_modulo_bitwidth_vec(
-; CHECK-NEXT:    [[R:%.*]] = call <2 x i32> @llvm.fshr.v2i32(<2 x i32> [[X:%.*]], <2 x i32> [[Y:%.*]], <2 x i32> <i32 34, i32 -1>)
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i32> @llvm.fshr.v2i32(<2 x i32> [[X:%.*]], <2 x i32> [[Y:%.*]], <2 x i32> <i32 2, i32 31>)
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %r = call <2 x i32> @llvm.fshr.v2i32(<2 x i32> %x, <2 x i32> %y, <2 x i32> <i32 34, i32 -1>)
   ret <2 x i32> %r
 }
 
 define <2 x i31> @fshl_constant_shift_amount_modulo_bitwidth_vec(<2 x i31> %x, <2 x i31> %y) {
 ; CHECK-LABEL: @fshl_constant_shift_amount_modulo_bitwidth_vec(

 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %r = call <2 x i32> @llvm.fshr.v2i32(<2 x i32> %x, <2 x i32> %y, <2 x i32> <i32 34, i32 -1>)
   ret <2 x i32> %r
 }
 
 define <2 x i31> @fshl_constant_shift_amount_modulo_bitwidth_vec(<2 x i31> %x, <2 x i31> %y) {
 ; CHECK-LABEL: @fshl_constant_shift_amount_modulo_bitwidth_vec(
-; CHECK-NEXT:    [[R:%.*]] = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> [[X:%.*]], <2 x i31> [[Y:%.*]], <2 x i31> <i31 34, i31 -1>)
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> [[X:%.*]], <2 x i31> [[Y:%.*]], <2 x i31> <i31 3, i31 1>)
 ; CHECK-NEXT:    ret <2 x i31> [[R]]
 ;
   %r = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> <i31 34, i31 -1>)
   ret <2 x i31> %r
 }
 
-; The shift modulo bitwidth is the same for all vector elements, but this is not simplified yet.
+define <2 x i31> @fshl_constant_shift_amount_modulo_bitwidth_vec_const_expr(<2 x i31> %x, <2 x i31> %y) {
+; CHECK-LABEL: @fshl_constant_shift_amount_modulo_bitwidth_vec_const_expr(
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> [[X:%.*]], <2 x i31> [[Y:%.*]], <2 x i31> <i31 34, i31 ptrtoint (i8* @external_global to i31)>)
+; CHECK-NEXT:    ret <2 x i31> [[R]]
+;
+  %shamt = ptrtoint i8* @external_global to i31
+  %r = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> <i31 34, i31 ptrtoint (i8* @external_global to i31)>)
+  ret <2 x i31> %r
+}
+
+; The shift modulo bitwidth is the same for all vector elements.
+
 define <2 x i31> @fshl_only_op1_demanded_vec_nonsplat(<2 x i31> %x, <2 x i31> %y) {
 ; CHECK-LABEL: @fshl_only_op1_demanded_vec_nonsplat(
-; CHECK-NEXT:    [[Z:%.*]] = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> [[X:%.*]], <2 x i31> [[Y:%.*]], <2 x i31> <i31 7, i31 38>)
+; CHECK-NEXT:    [[Z:%.*]] = lshr <2 x i31> [[Y:%.*]], <i31 24, i31 24>
 ; CHECK-NEXT:    [[R:%.*]] = and <2 x i31> [[Z]], <i31 63, i31 31>
 ; CHECK-NEXT:    ret <2 x i31> [[R]]
 ;
   %z = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> <i31 7, i31 38>)
   %r = and <2 x i31> %z, <i31 63, i31 31>
   ret <2 x i31> %r
 }