[InstCombine] visitTrunc - trunc (lshr (sext A), C) --> (ashr A, C) non-uniform support

[RKSimon]

Summary

This came from @lebedev.ri's suggestion to use m_SpecificInt_ICMP for D88429 - since I was going to change the m_APInt to m_Constant for that patch I thought I would do it for the only other user of the APInt first.

I've added a ConstantExpr::getUMin helper - its trivial to add UMAX/SMIN/SMAX but thought I'd wait until we have use cases.

[lebedev.ri]

Hm, why do we need that clamping here?
If it's for undefs, see Constant::replaceUndefsWith().

[RKSimon]

In D88475#2300260, @lebedev.ri wrote:

Hm, why do we need that clamping here?
If it's for undefs, see Constant::replaceUndefsWith().

Isn't it matching the existing std::min calls ?

[lebedev.ri]

In D88475#2300263, @RKSimon wrote:

In D88475#2300260, @lebedev.ri wrote:

Hm, why do we need that clamping here?
If it's for undefs, see Constant::replaceUndefsWith().

Isn't it matching the existing std::min calls ?

Yes, obviously. I'm asking why it's there in the first place.

[lebedev.ri]

I think this should look closer to: (precommit that before this patch)

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

830–857

[RKSimon]

Your patch breaks the trunc_lshr_sext_wide_input test resulting in this fail: https://alive2.llvm.org/ce/z/VmcSVq

If you don't mind I'd prefer to use my existing version instead of getting waylaid on this

[lebedev.ri]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

830–857

Err, that should of course be

// Note that since we will perform the shift in narrower type, we need to clamp the shift amount.
Constant *MaxAmt = ConstantInt::get(SrcTy, AWidth - 1, false);
Constant *ShAmt = ConstantExpr::getUMin(C, MaxAmt);
ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());

llvm/test/Transforms/InstCombine/cast.ll

1563

It would be kind of good to preserve undef, but there is no nice way to do that currently.

[lebedev.ri]

Note: here we also should be preserving exactness of the right-shift..

[lebedev.ri]

In D88475#2300561, @lebedev.ri wrote:

Note: here we also should be preserving exactness of the right-shift..

(i.e. unless you want to fix that in a followup, please can you add a FIXME comment there?)

[RKSimon]

In D88475#2300663, @lebedev.ri wrote:

In D88475#2300561, @lebedev.ri wrote:

Note: here we also should be preserving exactness of the right-shift..

(i.e. unless you want to fix that in a followup, please can you add a FIXME comment there?)

Sorry - missed your exactness comment - dealing with it now

llvm/include/llvm/IR/Constants.h

public:
   static Constant *getSDiv(Constant *C1, Constant *C2, bool isExact = false);
   static Constant *getFDiv(Constant *C1, Constant *C2);
   static Constant *getURem(Constant *C1, Constant *C2);
   static Constant *getSRem(Constant *C1, Constant *C2);
   static Constant *getFRem(Constant *C1, Constant *C2);
   static Constant *getAnd(Constant *C1, Constant *C2);
   static Constant *getOr(Constant *C1, Constant *C2);
   static Constant *getXor(Constant *C1, Constant *C2);
+  static Constant *getUMin(Constant *C1, Constant *C2);
   static Constant *getShl(Constant *C1, Constant *C2,
                           bool HasNUW = false, bool HasNSW = false);
   static Constant *getLShr(Constant *C1, Constant *C2, bool isExact = false);
   static Constant *getAShr(Constant *C1, Constant *C2, bool isExact = false);
   static Constant *getTrunc(Constant *C, Type *Ty, bool OnlyIfReduced = false);
   static Constant *getSExt(Constant *C, Type *Ty, bool OnlyIfReduced = false);
   static Constant *getZExt(Constant *C, Type *Ty, bool OnlyIfReduced = false);
   static Constant *getFPTrunc(Constant *C, Type *Ty,

llvm/lib/IR/Constants.cpp

 Constant *ConstantExpr::getOr(Constant *C1, Constant *C2) {
   return get(Instruction::Or, C1, C2);
 }
 
 Constant *ConstantExpr::getXor(Constant *C1, Constant *C2) {
   return get(Instruction::Xor, C1, C2);
 }
 
+Constant *ConstantExpr::getUMin(Constant *C1, Constant *C2) {
+  Constant *Cmp = ConstantExpr::getICmp(CmpInst::ICMP_ULT, C1, C2);
+  return getSelect(Cmp, C1, C2);
+}
+
 Constant *ConstantExpr::getShl(Constant *C1, Constant *C2,
                                bool HasNUW, bool HasNSW) {
   unsigned Flags = (HasNUW ? OverflowingBinaryOperator::NoUnsignedWrap : 0) |
                    (HasNSW ? OverflowingBinaryOperator::NoSignedWrap   : 0);
   return get(Instruction::Shl, C1, C2, Flags);
 }
 
 Constant *ConstantExpr::getLShr(Constant *C1, Constant *C2, bool isExact) {

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (Src->hasOneUse() &&
     // Since we're doing an lshr and a zero extend, and know that the shift
     // amount is smaller than ASize, it is always safe to do the shift in A's
     // type, then zero extend or truncate to the result.
     Value *Shift = Builder.CreateLShr(A, Cst->getZExtValue());
     Shift->takeName(Src);
     return CastInst::CreateIntegerCast(Shift, DestTy, false);
   }
 
-  const APInt *C;
-  if (match(Src, m_LShr(m_SExt(m_Value(A)), m_APInt(C)))) {
+  Constant *C;
+  if (match(Src, m_LShr(m_SExt(m_Value(A)), m_Constant(C)))) {
     unsigned AWidth = A->getType()->getScalarSizeInBits();
     unsigned MaxShiftAmt = SrcWidth - std::max(DestWidth, AWidth);
 
     // If the shift is small enough, all zero bits created by the shift are
     // removed by the trunc.
-    if (C->getZExtValue() <= MaxShiftAmt) {
+    // TODO: Support passing through undef shift amounts - these currently get
+    // clamped to MaxAmt.
+    if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
+                                    APInt(SrcWidth, MaxShiftAmt)))) {
       // trunc (lshr (sext A), C) --> ashr A, C
       if (A->getType() == DestTy) {
-        unsigned ShAmt = std::min((unsigned)C->getZExtValue(), DestWidth - 1);
-        return BinaryOperator::CreateAShr(A, ConstantInt::get(DestTy, ShAmt));
+        Constant *MaxAmt = ConstantInt::get(SrcTy, DestWidth - 1, false);
+        Constant *ShAmt = ConstantExpr::getUMin(C, MaxAmt);
+        ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());
+        return BinaryOperator::CreateAShr(A, ShAmt);
       }
       // The types are mismatched, so create a cast after shifting:
       // trunc (lshr (sext A), C) --> sext/trunc (ashr A, C)
       if (Src->hasOneUse()) {
-        unsigned ShAmt = std::min((unsigned)C->getZExtValue(), AWidth - 1);
+        Constant *MaxAmt = ConstantInt::get(SrcTy, AWidth - 1, false);
+        Constant *ShAmt = ConstantExpr::getUMin(C, MaxAmt);
+        ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());
         Value *Shift = Builder.CreateAShr(A, ShAmt);
         return CastInst::CreateIntegerCast(Shift, DestTy, true);
       }
     }
     // TODO: Mask high bits with 'and'.
   }
 
   if (Instruction *I = narrowBinOp(Trunc))
     return I;
 
   if (Instruction *I = shrinkSplatShuffle(Trunc, Builder))
     return I;

llvm/test/Transforms/InstCombine/cast.ll

;
   %B = sext <2 x i8> %A to <2 x i32>
   %C = lshr <2 x i32> %B, <i32 6, i32 6>
   %D = trunc <2 x i32> %C to <2 x i8>
   ret <2 x i8> %D
 }
 
 define <2 x i8> @trunc_lshr_sext_uniform_undef(<2 x i8> %A) {
 ; ALL-LABEL: @trunc_lshr_sext_uniform_undef(
-; ALL-NEXT:    [[B:%.*]] = sext <2 x i8> [[A:%.*]] to <2 x i32>
-; ALL-NEXT:    [[C:%.*]] = lshr <2 x i32> [[B]], <i32 6, i32 undef>
-; ALL-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; ALL-NEXT:    [[D:%.*]] = ashr <2 x i8> [[A:%.*]], <i8 6, i8 7>
 ; ALL-NEXT:    ret <2 x i8> [[D]]
 ;
   %B = sext <2 x i8> %A to <2 x i32>
   %C = lshr <2 x i32> %B, <i32 6, i32 undef>
   %D = trunc <2 x i32> %C to <2 x i8>
   ret <2 x i8> %D
 }
 
 define <2 x i8> @trunc_lshr_sext_nonuniform(<2 x i8> %A) {
 ; ALL-LABEL: @trunc_lshr_sext_nonuniform(
-; ALL-NEXT:    [[B:%.*]] = sext <2 x i8> [[A:%.*]] to <2 x i32>
-; ALL-NEXT:    [[C:%.*]] = lshr <2 x i32> [[B]], <i32 6, i32 2>
-; ALL-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; ALL-NEXT:    [[D:%.*]] = ashr <2 x i8> [[A:%.*]], <i8 6, i8 2>
 ; ALL-NEXT:    ret <2 x i8> [[D]]
 ;
   %B = sext <2 x i8> %A to <2 x i32>
   %C = lshr <2 x i32> %B, <i32 6, i32 2>
   %D = trunc <2 x i32> %C to <2 x i8>
   ret <2 x i8> %D
 }
 
 define <3 x i8> @trunc_lshr_sext_nonuniform_undef(<3 x i8> %A) {
 ; ALL-LABEL: @trunc_lshr_sext_nonuniform_undef(
-; ALL-NEXT:    [[B:%.*]] = sext <3 x i8> [[A:%.*]] to <3 x i32>
-; ALL-NEXT:    [[C:%.*]] = lshr <3 x i32> [[B]], <i32 6, i32 2, i32 undef>
-; ALL-NEXT:    [[D:%.*]] = trunc <3 x i32> [[C]] to <3 x i8>
+; ALL-NEXT:    [[D:%.*]] = ashr <3 x i8> [[A:%.*]], <i8 6, i8 2, i8 7>
 ; ALL-NEXT:    ret <3 x i8> [[D]]
 ;
   %B = sext <3 x i8> %A to <3 x i32>
   %C = lshr <3 x i32> %B, <i32 6, i32 2, i32 undef>
   %D = trunc <3 x i32> %C to <3 x i8>
   ret <3 x i8> %D
 }