[InstCombine] visitTrunc - trunc(shl(X, C)) --> shl(trunc(X),trunc(C)) vector support

[RKSimon]

Summary

Annoyingly vectors aren't supported by shouldChangeType(), but we have precedents for always performing this on vector types (e.g. narrowBinOp)

[lebedev.ri]

Seems rather straight-forward to me, thanks.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

 Instruction *InstCombinerImpl::visitTrunc(TruncInst &Trunc) {
   if (Instruction *Result = commonCastTransforms(Trunc))
     return Result;
 
   Value *Src = Trunc.getOperand(0);
   Type *DestTy = Trunc.getType(), *SrcTy = Src->getType();
   unsigned DestWidth = DestTy->getScalarSizeInBits();
   unsigned SrcWidth = SrcTy->getScalarSizeInBits();
-  ConstantInt *Cst;
 
   // Attempt to truncate the entire input expression tree to the destination
   // type.   Only do this if the dest type is a simple type, don't convert the
   // expression tree to something weird like i93 unless the source is also
   // strange.
   if ((DestTy->isVectorTy() || shouldChangeType(SrcTy, DestTy)) &&
       canEvaluateTruncated(Src, DestTy, *this, &Trunc)) {
 
if (Instruction *I = narrowBinOp(Trunc))
     return I;
 
   if (Instruction *I = shrinkSplatShuffle(Trunc, Builder))
     return I;
 
   if (Instruction *I = shrinkInsertElt(Trunc, Builder))
     return I;
 
-  if (Src->hasOneUse() && isa<IntegerType>(SrcTy) &&
-      shouldChangeType(SrcTy, DestTy)) {
+  if (Src->hasOneUse() &&
+      (isa<VectorType>(SrcTy) || shouldChangeType(SrcTy, DestTy))) {
     // Transform "trunc (shl X, cst)" -> "shl (trunc X), cst" so long as the
     // dest type is native and cst < dest size.
-    if (match(Src, m_Shl(m_Value(A), m_ConstantInt(Cst))) &&
+    if (match(Src, m_Shl(m_Value(A), m_Constant(C))) &&
         !match(A, m_Shr(m_Value(), m_Constant()))) {
       // Skip shifts of shift by constants. It undoes a combine in
       // FoldShiftByConstant and is the extend in reg pattern.
-      if (Cst->getValue().ult(DestWidth)) {
+      APInt Threshold = APInt(C->getType()->getScalarSizeInBits(), DestWidth);
+      if (match(C, m_SpecificInt_ICMP(ICmpInst::ICMP_ULT, Threshold))) {
         Value *NewTrunc = Builder.CreateTrunc(A, DestTy, A->getName() + ".tr");
-
-        return BinaryOperator::Create(
-          Instruction::Shl, NewTrunc,
-          ConstantInt::get(DestTy, Cst->getValue().trunc(DestWidth)));
+        return BinaryOperator::Create(Instruction::Shl, NewTrunc,
+                                      ConstantExpr::getTrunc(C, DestTy));
       }
     }
   }
 
   if (Instruction *I = foldVecTruncToExtElt(Trunc, *this))
     return I;
 
   // Whenever an element is extracted from a vector, and then truncated,
   // canonicalize by converting it to a bitcast followed by an
   // extractelement.
   //
   // Example (little endian):
   //   trunc (extractelement <4 x i64> %X, 0) to i32
   //   --->
   //   extractelement <8 x i32> (bitcast <4 x i64> %X to <8 x i32>), i32 0
   Value *VecOp;
+  ConstantInt *Cst;
   if (match(Src, m_OneUse(m_ExtractElt(m_Value(VecOp), m_ConstantInt(Cst))))) {
     auto *VecOpTy = cast<FixedVectorType>(VecOp->getType());
     unsigned VecNumElts = VecOpTy->getNumElements();
 
     // A badly fit destination size would result in an invalid cast.
     if (SrcWidth % DestWidth == 0) {
       uint64_t TruncRatio = SrcWidth / DestWidth;
       uint64_t BitCastNumElts = VecNumElts * TruncRatio;

llvm/test/Transforms/InstCombine/shift-amount-reassociation-in-bittest-with-truncation-shl.ll

;
   %t1_trunc = trunc i64 %t1 to i32
   %t2 = and i32 %t0, %t1_trunc
   %t3 = icmp ne i32 %t2, 0
   ret i1 %t3
 }
 
 define <2 x i1> @t11_constants_vec_splat(<2 x i32> %x, <2 x i64> %y) {
 ; CHECK-LABEL: @t11_constants_vec_splat(
-; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 12, i32 12>
-; CHECK-NEXT:    [[T1:%.*]] = shl <2 x i64> [[Y:%.*]], <i64 14, i64 14>
-; CHECK-NEXT:    [[T1_TRUNC:%.*]] = trunc <2 x i64> [[T1]] to <2 x i32>
-; CHECK-NEXT:    [[T2:%.*]] = and <2 x i32> [[T0]], [[T1_TRUNC]]
-; CHECK-NEXT:    [[T3:%.*]] = icmp ne <2 x i32> [[T2]], zeroinitializer
-; CHECK-NEXT:    ret <2 x i1> [[T3]]
+; CHECK-NEXT:    [[Y_TR:%.*]] = trunc <2 x i64> [[Y:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 26, i32 26>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], [[Y_TR]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne <2 x i32> [[TMP2]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[TMP3]]
 ;
   %t0 = lshr <2 x i32> %x, <i32 12, i32 12>
   %t1 = shl <2 x i64> %y, <i64 14, i64 14>
   %t1_trunc = trunc <2 x i64> %t1 to <2 x i32>
   %t2 = and <2 x i32> %t0, %t1_trunc
   %t3 = icmp ne <2 x i32> %t2, <i32 0, i32 0>
   ret <2 x i1> %t3
 }
 define <2 x i1> @t12_constants_vec_nonsplat(<2 x i32> %x, <2 x i64> %y) {
 ; CHECK-LABEL: @t12_constants_vec_nonsplat(
-; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 12, i32 14>
-; CHECK-NEXT:    [[T1:%.*]] = shl <2 x i64> [[Y:%.*]], <i64 16, i64 14>
-; CHECK-NEXT:    [[T1_TRUNC:%.*]] = trunc <2 x i64> [[T1]] to <2 x i32>
-; CHECK-NEXT:    [[T2:%.*]] = and <2 x i32> [[T0]], [[T1_TRUNC]]
-; CHECK-NEXT:    [[T3:%.*]] = icmp ne <2 x i32> [[T2]], zeroinitializer
-; CHECK-NEXT:    ret <2 x i1> [[T3]]
+; CHECK-NEXT:    [[Y_TR:%.*]] = trunc <2 x i64> [[Y:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 28, i32 28>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], [[Y_TR]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne <2 x i32> [[TMP2]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[TMP3]]
 ;
   %t0 = lshr <2 x i32> %x, <i32 12, i32 14>
   %t1 = shl <2 x i64> %y, <i64 16, i64 14>
   %t1_trunc = trunc <2 x i64> %t1 to <2 x i32>
   %t2 = and <2 x i32> %t0, %t1_trunc
   %t3 = icmp ne <2 x i32> %t2, <i32 0, i32 0>
   ret <2 x i1> %t3
 }

llvm/test/Transforms/InstCombine/shift-amount-reassociation-with-truncation-shl.ll

;
   %t3 = trunc i32 %t2 to i16
   %t4 = add i16 %y, -24
   %t5 = shl i16 %t3, %t4
   ret i16 %t5
 }
 
 define <2 x i16> @t1_vec_splat(<2 x i32> %x, <2 x i16> %y) {
 ; CHECK-LABEL: @t1_vec_splat(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl <2 x i32> [[X:%.*]], <i32 8, i32 8>
-; CHECK-NEXT:    [[T5:%.*]] = trunc <2 x i32> [[TMP1]] to <2 x i16>
+; CHECK-NEXT:    [[X_TR:%.*]] = trunc <2 x i32> [[X:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[T5:%.*]] = shl <2 x i16> [[X_TR]], <i16 8, i16 8>
 ; CHECK-NEXT:    ret <2 x i16> [[T5]]
 ;
   %t0 = sub <2 x i16> <i16 32, i16 32>, %y
   %t1 = zext <2 x i16> %t0 to <2 x i32>
   %t2 = shl <2 x i32> %x, %t1
   %t3 = trunc <2 x i32> %t2 to <2 x i16>
   %t4 = add <2 x i16> %y, <i16 -24, i16 -24>
   %t5 = shl <2 x i16> %t3, %t4
;
   %t5 = shl <2 x i16> %t3, %t4
   ret <2 x i16> %t5
 }
 
 ; Basic vector tests
 
 define <3 x i16> @t3_vec_nonsplat_undef0(<3 x i32> %x, <3 x i16> %y) {
 ; CHECK-LABEL: @t3_vec_nonsplat_undef0(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl <3 x i32> [[X:%.*]], <i32 8, i32 0, i32 8>
-; CHECK-NEXT:    [[T5:%.*]] = trunc <3 x i32> [[TMP1]] to <3 x i16>
+; CHECK-NEXT:    [[X_TR:%.*]] = trunc <3 x i32> [[X:%.*]] to <3 x i16>
+; CHECK-NEXT:    [[T5:%.*]] = shl <3 x i16> [[X_TR]], <i16 8, i16 0, i16 8>
 ; CHECK-NEXT:    ret <3 x i16> [[T5]]
 ;
   %t0 = sub <3 x i16> <i16 32, i16 undef, i16 32>, %y
   %t1 = zext <3 x i16> %t0 to <3 x i32>
   %t2 = shl <3 x i32> %x, %t1
   %t3 = trunc <3 x i32> %t2 to <3 x i16>
   %t4 = add <3 x i16> %y, <i16 -24, i16 -24, i16 -24>
   %t5 = shl <3 x i16> %t3, %t4
   ret <3 x i16> %t5
 }
 
 define <3 x i16> @t4_vec_nonsplat_undef1(<3 x i32> %x, <3 x i16> %y) {
 ; CHECK-LABEL: @t4_vec_nonsplat_undef1(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl <3 x i32> [[X:%.*]], <i32 8, i32 0, i32 8>
-; CHECK-NEXT:    [[T5:%.*]] = trunc <3 x i32> [[TMP1]] to <3 x i16>
+; CHECK-NEXT:    [[X_TR:%.*]] = trunc <3 x i32> [[X:%.*]] to <3 x i16>
+; CHECK-NEXT:    [[T5:%.*]] = shl <3 x i16> [[X_TR]], <i16 8, i16 0, i16 8>
 ; CHECK-NEXT:    ret <3 x i16> [[T5]]
 ;
   %t0 = sub <3 x i16> <i16 32, i16 32, i16 32>, %y
   %t1 = zext <3 x i16> %t0 to <3 x i32>
   %t2 = shl <3 x i32> %x, %t1
   %t3 = trunc <3 x i32> %t2 to <3 x i16>
   %t4 = add <3 x i16> %y, <i16 -24, i16 undef, i16 -24>
   %t5 = shl <3 x i16> %t3, %t4
   ret <3 x i16> %t5
 }
 
 define <3 x i16> @t5_vec_nonsplat_undef1(<3 x i32> %x, <3 x i16> %y) {
 ; CHECK-LABEL: @t5_vec_nonsplat_undef1(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl <3 x i32> [[X:%.*]], <i32 8, i32 0, i32 8>
-; CHECK-NEXT:    [[T5:%.*]] = trunc <3 x i32> [[TMP1]] to <3 x i16>
+; CHECK-NEXT:    [[X_TR:%.*]] = trunc <3 x i32> [[X:%.*]] to <3 x i16>
+; CHECK-NEXT:    [[T5:%.*]] = shl <3 x i16> [[X_TR]], <i16 8, i16 0, i16 8>
 ; CHECK-NEXT:    ret <3 x i16> [[T5]]
 ;
   %t0 = sub <3 x i16> <i16 32, i16 undef, i16 32>, %y
   %t1 = zext <3 x i16> %t0 to <3 x i32>
   %t2 = shl <3 x i32> %x, %t1
   %t3 = trunc <3 x i32> %t2 to <3 x i16>
   %t4 = add <3 x i16> %y, <i16 -24, i16 undef, i16 -24>
   %t5 = shl <3 x i16> %t3, %t4

llvm/test/Transforms/InstCombine/trunc.ll

 ; CHECK-LABEL: @trunc_shl_32_i32_i64(
 ; CHECK-NEXT:    ret i32 0
 ;
   %shl = shl i64 %val, 32
   %trunc = trunc i64 %shl to i32
   ret i32 %trunc
 }
 
-; TODO: Should be able to handle vectors
+; Should be able to handle vectors
 define <2 x i32> @trunc_shl_16_v2i32_v2i64(<2 x i64> %val) {
 ; CHECK-LABEL: @trunc_shl_16_v2i32_v2i64(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i64> [[VAL:%.*]], <i64 16, i64 16>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i64> [[SHL]] to <2 x i32>
+; CHECK-NEXT:    [[VAL_TR:%.*]] = trunc <2 x i64> [[VAL:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shl <2 x i32> [[VAL_TR]], <i32 16, i32 16>
 ; CHECK-NEXT:    ret <2 x i32> [[TRUNC]]
 ;
   %shl = shl <2 x i64> %val, <i64 16, i64 16>
   %trunc = trunc <2 x i64> %shl to <2 x i32>
   ret <2 x i32> %trunc
 }
 
 define <2 x i32> @trunc_shl_nosplat_v2i32_v2i64(<2 x i64> %val) {
 ; CHECK-LABEL: @trunc_shl_nosplat_v2i32_v2i64(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i64> [[VAL:%.*]], <i64 15, i64 16>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i64> [[SHL]] to <2 x i32>
+; CHECK-NEXT:    [[VAL_TR:%.*]] = trunc <2 x i64> [[VAL:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shl <2 x i32> [[VAL_TR]], <i32 15, i32 16>
 ; CHECK-NEXT:    ret <2 x i32> [[TRUNC]]
 ;
   %shl = shl <2 x i64> %val, <i64 15, i64 16>
   %trunc = trunc <2 x i64> %shl to <2 x i32>
   ret <2 x i32> %trunc
 }
 
 define void @trunc_shl_31_i32_i64_multi_use(i64 %val, i32 addrspace(1)* %ptr0, i64 addrspace(1)* %ptr1) {
;
   %A = lshr i64 %arg, 1
   %B = shl i64 %A, 2
   %C = trunc i64 %B to i32
   ret i32 %C
 }
 
 define <2 x i32> @trunc_shl_v2i32_v2i64_uniform(<2 x i64> %val) {
 ; CHECK-LABEL: @trunc_shl_v2i32_v2i64_uniform(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i64> [[VAL:%.*]], <i64 31, i64 31>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i64> [[SHL]] to <2 x i32>
+; CHECK-NEXT:    [[VAL_TR:%.*]] = trunc <2 x i64> [[VAL:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shl <2 x i32> [[VAL_TR]], <i32 31, i32 31>
 ; CHECK-NEXT:    ret <2 x i32> [[TRUNC]]
 ;
   %shl = shl <2 x i64> %val, <i64 31, i64 31>
   %trunc = trunc <2 x i64> %shl to <2 x i32>
   ret <2 x i32> %trunc
 }
 
 define <2 x i32> @trunc_shl_v2i32_v2i64_undef(<2 x i64> %val) {
 ; CHECK-LABEL: @trunc_shl_v2i32_v2i64_undef(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i64> [[VAL:%.*]], <i64 31, i64 undef>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i64> [[SHL]] to <2 x i32>
+; CHECK-NEXT:    [[VAL_TR:%.*]] = trunc <2 x i64> [[VAL:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shl <2 x i32> [[VAL_TR]], <i32 31, i32 undef>
 ; CHECK-NEXT:    ret <2 x i32> [[TRUNC]]
 ;
   %shl = shl <2 x i64> %val, <i64 31, i64 undef>
   %trunc = trunc <2 x i64> %shl to <2 x i32>
   ret <2 x i32> %trunc
 }
 
 define <2 x i32> @trunc_shl_v2i32_v2i64_nonuniform(<2 x i64> %val) {
 ; CHECK-LABEL: @trunc_shl_v2i32_v2i64_nonuniform(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i64> [[VAL:%.*]], <i64 31, i64 12>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i64> [[SHL]] to <2 x i32>
+; CHECK-NEXT:    [[VAL_TR:%.*]] = trunc <2 x i64> [[VAL:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shl <2 x i32> [[VAL_TR]], <i32 31, i32 12>
 ; CHECK-NEXT:    ret <2 x i32> [[TRUNC]]
 ;
   %shl = shl <2 x i64> %val, <i64 31, i64 12>
   %trunc = trunc <2 x i64> %shl to <2 x i32>
   ret <2 x i32> %trunc
 }
 
 define <2 x i32> @trunc_shl_v2i32_v2i64_outofrange(<2 x i64> %val) {
;
   %A = shl i64 %arg, %val
   %B = shl i64 %A, 2
   %C = trunc i64 %B to i32
   ret i32 %C
 }
 
 define <8 x i16> @trunc_shl_v8i15_v8i32_15(<8 x i32> %a) {
 ; CHECK-LABEL: @trunc_shl_v8i15_v8i32_15(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <8 x i32> [[A:%.*]], <i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15>
-; CHECK-NEXT:    [[CONV:%.*]] = trunc <8 x i32> [[SHL]] to <8 x i16>
+; CHECK-NEXT:    [[A_TR:%.*]] = trunc <8 x i32> [[A:%.*]] to <8 x i16>
+; CHECK-NEXT:    [[CONV:%.*]] = shl <8 x i16> [[A_TR]], <i16 15, i16 15, i16 15, i16 15, i16 15, i16 15, i16 15, i16 15>
 ; CHECK-NEXT:    ret <8 x i16> [[CONV]]
 ;
   %shl = shl <8 x i32> %a, <i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15>
   %conv = trunc <8 x i32> %shl to <8 x i16>
   ret <8 x i16> %conv
 }
 
 define <8 x i16> @trunc_shl_v8i16_v8i32_16(<8 x i32> %a) {

 ;
   %shl = shl <8 x i32> %a, <i32 17, i32 17, i32 17, i32 17, i32 17, i32 17, i32 17, i32 17>
   %conv = trunc <8 x i32> %shl to <8 x i16>
   ret <8 x i16> %conv
 }
 
 define <8 x i16> @trunc_shl_v8i16_v8i32_4(<8 x i32> %a) {
 ; CHECK-LABEL: @trunc_shl_v8i16_v8i32_4(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <8 x i32> [[A:%.*]], <i32 4, i32 4, i32 4, i32 4, i32 4, i32 4, i32 4, i32 4>
-; CHECK-NEXT:    [[CONV:%.*]] = trunc <8 x i32> [[SHL]] to <8 x i16>
+; CHECK-NEXT:    [[A_TR:%.*]] = trunc <8 x i32> [[A:%.*]] to <8 x i16>
+; CHECK-NEXT:    [[CONV:%.*]] = shl <8 x i16> [[A_TR]], <i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4>
 ; CHECK-NEXT:    ret <8 x i16> [[CONV]]
 ;
   %shl = shl <8 x i32> %a, <i32 4, i32 4, i32 4, i32 4, i32 4, i32 4, i32 4, i32 4>
   %conv = trunc <8 x i32> %shl to <8 x i16>
   ret <8 x i16> %conv
 }
 
 ; Although the mask is the same value, we don't create a shuffle for types that the backend may not be able to handle: