[InstCombine] [NFC] Move transforms for truncated shifts into narrowBinOp

[spatel]

Please make sure there is good test coverage and no compile-time problems before posting a patch for review.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

687–710

This can't be correct - we are replacing shl with lshr?

709

This case falls through and creates a compile-time warning (when compiled by clang).

[Chenbing.Zheng]

address comments

[spatel]

LGTM - but please add a test to verify that we do not create a bug with shl.

I think this would have miscompiled with the previous revision of the patch:

define <2 x i8> @trunc_lshr_trunc(<2 x i64> %a) {
  %b = trunc <2 x i64> %a to <2 x i32>
  %c = shl <2 x i32> %b, <i32 9, i32 7>
  %d = trunc <2 x i32> %c to <2 x i8>
  ret <2 x i8> %d
}

[Chenbing.Zheng]

In D126056#3534059, @spatel wrote:

LGTM - but please add a test to verify that we do not create a bug with shl.

I think this would have miscompiled with the previous revision of the patch:

define <2 x i8> @trunc_lshr_trunc(<2 x i64> %a) {
  %b = trunc <2 x i64> %a to <2 x i32>
  %c = shl <2 x i32> %b, <i32 9, i32 7>
  %d = trunc <2 x i32> %c to <2 x i8>
  ret <2 x i8> %d
}

I commit it in https://reviews.llvm.org/rG793bb7049db012df4fe3ff6e8abc320677845400，and it pass after this patch.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (ShVal0 != ShVal1)
     Y = Builder.CreateTrunc(ShVal1, DestTy);
   Intrinsic::ID IID = IsFshl ? Intrinsic::fshl : Intrinsic::fshr;
   Function *F = Intrinsic::getDeclaration(Trunc.getModule(), IID, DestTy);
   return CallInst::Create(F, {X, Y, NarrowShAmt});
 }
 
 /// Try to narrow the width of math or bitwise logic instructions by pulling a
 /// truncate ahead of binary operators.
-/// TODO: Transforms for truncated shifts should be moved into here.
 Instruction *InstCombinerImpl::narrowBinOp(TruncInst &Trunc) {
   Type *SrcTy = Trunc.getSrcTy();
   Type *DestTy = Trunc.getType();
+  unsigned SrcWidth = SrcTy->getScalarSizeInBits();
+  unsigned DestWidth = DestTy->getScalarSizeInBits();
+
   if (!isa<VectorType>(SrcTy) && !shouldChangeType(SrcTy, DestTy))
     return nullptr;
 
   BinaryOperator *BinOp;
   if (!match(Trunc.getOperand(0), m_OneUse(m_BinOp(BinOp))))
     return nullptr;
 
   Value *BinOp0 = BinOp->getOperand(0);
case Instruction::Mul: {
     }
     if (match(BinOp1, m_ZExtOrSExt(m_Value(X))) && X->getType() == DestTy) {
       // trunc (binop Y, (ext X)) --> binop (trunc Y), X
       Value *NarrowOp0 = Builder.CreateTrunc(BinOp0, DestTy);
       return BinaryOperator::Create(BinOp->getOpcode(), NarrowOp0, X);
     }
     break;
   }
-
+  case Instruction::LShr:
+  case Instruction::AShr: {
+    // trunc (*shr (trunc A), C) --> trunc(*shr A, C)
+    Value *A;
+    Constant *C;
+    if (match(BinOp0, m_Trunc(m_Value(A))) && match(BinOp1, m_Constant(C))) {
+      unsigned MaxShiftAmt = SrcWidth - DestWidth;
+      // If the shift is small enough, all zero/sign bits created by the shift
+      // are removed by the trunc.
+      if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
+                                      APInt(SrcWidth, MaxShiftAmt)))) {
+        auto *OldShift = cast<Instruction>(Trunc.getOperand(0));
+        bool IsExact = OldShift->isExact();
+        auto *ShAmt = ConstantExpr::getIntegerCast(C, A->getType(), true);
+        ShAmt = Constant::mergeUndefsWith(ShAmt, C);
+        Value *Shift =
+            OldShift->getOpcode() == Instruction::AShr
+                ? Builder.CreateAShr(A, ShAmt, OldShift->getName(), IsExact)
+                : Builder.CreateLShr(A, ShAmt, OldShift->getName(), IsExact);
+        return CastInst::CreateTruncOrBitCast(Shift, DestTy);
+      }
+    }
+    break;
+  }
   default: break;
   }
 
   if (Instruction *NarrowOr = narrowFunnelShift(Trunc))
     return NarrowOr;
 
   return nullptr;
 }
if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
         ShAmt = ConstantExpr::getTrunc(ShAmt, A->getType());
         Value *Shift = Builder.CreateAShr(A, ShAmt, "", IsExact);
         return CastInst::CreateIntegerCast(Shift, DestTy, true);
       }
     }
     // TODO: Mask high bits with 'and'.
   }
 
-  // trunc (*shr (trunc A), C) --> trunc(*shr A, C)
-  if (match(Src, m_OneUse(m_Shr(m_Trunc(m_Value(A)), m_Constant(C))))) {
-    unsigned MaxShiftAmt = SrcWidth - DestWidth;
-
-    // If the shift is small enough, all zero/sign bits created by the shift are
-    // removed by the trunc.
-    if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
-                                    APInt(SrcWidth, MaxShiftAmt)))) {
-      auto *OldShift = cast<Instruction>(Src);
-      bool IsExact = OldShift->isExact();
-      auto *ShAmt = ConstantExpr::getIntegerCast(C, A->getType(), true);
-      ShAmt = Constant::mergeUndefsWith(ShAmt, C);
-      Value *Shift =
-          OldShift->getOpcode() == Instruction::AShr
-              ? Builder.CreateAShr(A, ShAmt, OldShift->getName(), IsExact)
-              : Builder.CreateLShr(A, ShAmt, OldShift->getName(), IsExact);
-      return CastInst::CreateTruncOrBitCast(Shift, DestTy);
-    }
-  }
-
   if (Instruction *I = narrowBinOp(Trunc))
     return I;
 
   if (Instruction *I = shrinkSplatShuffle(Trunc, Builder))
     return I;
 
   if (Instruction *I = shrinkInsertElt(Trunc, Builder))
     return I;