[VectorCombine] Narrow ZExt that feed binop followed by trunc.

[fhahn]

Summary

In the pattern below, the trunc can be eliminated by shortening the
zexts, if the zexts remain.

trunc (binop (zext), (zext)) to ty -> binop (zext to ty) (zext to ty)

Initially limited to add/sub.

This transform is only performed if the shortened zexts are free (can be
folded into the binary op).

I am not entirely sure VectorCombine is the right place to do the
transform, but I think we want to limit it to cases where we know the
shorter zexts are free/legal on the target. I am not sure if we have an
easy way to check the latter though.

Alive proof sketches (scalar versions so we do not run into timeouts):

• add: https://alive2.llvm.org/ce/z/DgABb-
• add nuw: https://alive2.llvm.org/ce/z/yx5Vag
• add nsw: https://alive2.llvm.org/ce/z/yyVoRU
• sub: https://alive2.llvm.org/ce/z/bKj22_
• sub nuw: https://alive2.llvm.org/ce/z/H8soWR
• sub nsw: https://alive2.llvm.org/ce/z/QLVNDK

On AArch64, codegen for the following input can be improved (this is
from hot code in SPEC2006/h264)

define <8 x i32> @test(<8 x i16>* %p1, <8 x i16>* %p2) {

%l.1 = load <8 x i16>, <8 x i16>* %p1, align 2
%ext.1 = zext <8 x i16> %l.1 to <8 x i64>
%l.2 = load <8 x i16>, <8 x i16>* %p2, align 2
%ext.2 = zext <8 x i16> %l.2 to <8 x i64>
%sub = sub nsw <8 x i64> %ext.1, %ext.2
%t = trunc <8 x i64> %sub to <8 x i32>
ret <8 x i32> %t

}

Without patch

ldr     q0, [x0]
ldr     q1, [x1]
ushll2  v2.4s, v0.8h, #0
ushll   v0.4s, v0.4h, #0
ushll2  v3.4s, v1.8h, #0
ushll   v1.4s, v1.4h, #0
usubl2  v4.2d, v0.4s, v1.4s
usubl   v0.2d, v0.2s, v1.2s
usubl   v1.2d, v2.2s, v3.2s
usubl2  v5.2d, v2.4s, v3.4s
xtn     v1.2s, v1.2d
xtn     v0.2s, v0.2d
xtn2    v1.4s, v5.2d
xtn2    v0.4s, v4.2d
ret

With patch

ldr     q0, [x0]
ldr     q2, [x1]
usubl2  v1.4s, v0.8h, v2.8h
usubl   v0.4s, v0.4h, v2.4h
ret

[fhahn]

Move llvm/test/Transforms/VectorCombine/AArch64/lit.local.cfg to NFC test patch.

[RKSimon]

Doesn't instcombine handle this already? https://godbolt.org/z/TQLYd3

[fhahn]

In D83135#2130707, @RKSimon wrote:

Doesn't instcombine handle this already? https://godbolt.org/z/TQLYd3

Oh indeed, that's great. Looks like this patch is not needed :)

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

 
 #define DEBUG_TYPE "vector-combine"
 STATISTIC(NumVecCmp, "Number of vector compares formed");
 STATISTIC(NumVecBO, "Number of vector binops formed");
 STATISTIC(NumVecCmpBO, "Number of vector compare + binop formed");
 STATISTIC(NumShufOfBitcast, "Number of shuffles moved after bitcast");
 STATISTIC(NumScalarBO, "Number of scalar binops formed");
 STATISTIC(NumScalarCmp, "Number of scalar compares formed");
+STATISTIC(NumTruncElim, "Number of truncates eliminated");
 
 static cl::opt<bool> DisableVectorCombine(
     "disable-vector-combine", cl::init(false), cl::Hidden,
     cl::desc("Disable all vector combine transforms"));
 
 static cl::opt<bool> DisableBinopExtractShuffle(
     "disable-binop-extract-shuffle", cl::init(false), cl::Hidden,
     cl::desc("Disable binop extract to shuffle transforms"));
private:
   void foldExtExtCmp(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                      Instruction &I);
   void foldExtExtBinop(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                        Instruction &I);
   bool foldExtractExtract(Instruction &I);
   bool foldBitcastShuf(Instruction &I);
   bool scalarizeBinopOrCmp(Instruction &I);
   bool foldExtractedCmps(Instruction &I);
+  /// Try to shorten zexts so if it allows eliminating a trunc and the shorter
+  /// zexts are free.
+  bool foldTruncBinopExt(Instruction &I);
 };
 
 static void replaceValue(Value &Old, Value &New) {
   Old.replaceAllUsesWith(&New);
   New.takeName(&Old);
 }
 
 /// Determine which, if any, of the inputs should be replaced by a shuffle
bool VectorCombine::foldExtractedCmps(Instruction &I) {
   Value *VecLogic = Builder.CreateBinOp(cast<BinaryOperator>(I).getOpcode(),
                                         VCmp, Shuf);
   Value *NewExt = Builder.CreateExtractElement(VecLogic, CheapIndex);
   replaceValue(I, *NewExt);
   ++NumVecCmpBO;
   return true;
 }
 
+bool VectorCombine::foldTruncBinopExt(Instruction &I) {
+  Value *V1;
+  Value *V2;
+  if (!I.getType()->isVectorTy())
+    return false;
+
+  BinaryOperator *BO;
+  if (!match(&I, m_Trunc(m_OneUse(m_CombineAnd(
+                     m_BinOp(m_OneUse(m_Value(V1)), m_OneUse(m_Value(V2))),
+                     m_BinOp(BO))))))
+    return false;
+
+  // Limit transform to known-safe opcodes.
+  // TODO: Add more safe opcodes.
+  if (BO->getOpcode() != Instruction::Add &&
+      BO->getOpcode() != Instruction::Sub)
+    return false;
+
+  // Limit to ZExt operands that have the same input width and are at least as
+  // large as the truncate width.
+  auto *ZExt1 = dyn_cast<ZExtInst>(V1);
+  auto *ZExt2 = dyn_cast<ZExtInst>(V2);
+  if (!ZExt1 || !ZExt2)
+    return false;
+  if (ZExt1->getOperand(0)->getType() != ZExt2->getOperand(0)->getType())
+    return false;
+  if (I.getType()->getScalarSizeInBits() <=
+      ZExt1->getOperand(0)->getType()->getScalarSizeInBits())
+    return false;
+
+  // Check if the new ZExts can be done for free by folding them into their
+  // user.
+  if (TTI.getCastInstrCost(ZExt1->getOpcode(), I.getType(), V1->getType(),
+                           TTI::TCK_RecipThroughput, ZExt1) != 0)
+    return false;
+
+  auto *ShortZExt1 = Builder.CreateZExt(ZExt1->getOperand(0), I.getType());
+  auto *ShortZExt2 = Builder.CreateZExt(ZExt2->getOperand(0), I.getType());
+  auto *ShortSub = Builder.CreateBinOp(BO->getOpcode(), ShortZExt1, ShortZExt2);
+  auto *ShortBO = dyn_cast<BinaryOperator>(ShortSub);
+  if (BO->hasNoUnsignedWrap())
+    ShortBO->setHasNoUnsignedWrap();
+  if (BO->hasNoSignedWrap())
+    ShortBO->setHasNoSignedWrap();
+  NumTruncElim++;
+  I.replaceAllUsesWith(ShortSub);
+  return true;
+}
+
 /// This is the entry point for all transforms. Pass manager differences are
 /// handled in the callers of this function.
 bool VectorCombine::run() {
   if (DisableVectorCombine)
     return false;
 
   bool MadeChange = false;
   for (BasicBlock &BB : F) {
     // Ignore unreachable basic blocks.
     if (!DT.isReachableFromEntry(&BB))
       continue;
     // Do not delete instructions under here and invalidate the iterator.
     // Walk the block forwards to enable simple iterative chains of transforms.
     // TODO: It could be more efficient to remove dead instructions
     //       iteratively in this loop rather than waiting until the end.
     for (Instruction &I : BB) {
       if (isa<DbgInfoIntrinsic>(I))
         continue;
       Builder.SetInsertPoint(&I);
       MadeChange |= foldExtractExtract(I);
       MadeChange |= foldBitcastShuf(I);
       MadeChange |= scalarizeBinopOrCmp(I);
       MadeChange |= foldExtractedCmps(I);
+      MadeChange |= foldTruncBinopExt(I);
     }
   }
 
   // We're done with transforms, so remove dead instructions.
   if (MadeChange)
     for (BasicBlock &BB : F)
       SimplifyInstructionsInBlock(&BB);
 

llvm/test/Transforms/VectorCombine/AArch64/shorten-extend-if-free.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -vector-combine -mtriple=arm64-apple-iphoneos -S %s | FileCheck %s
 
 ; Zext can be folded with sub add for free.
 define <8 x i32> @shorten_zext_sub(<8 x i16> %v.1, <8 x i16> %v.2) {
 ; CHECK-LABEL: @shorten_zext_sub(
-; CHECK-NEXT:    [[EXT_1:%.*]] = zext <8 x i16> [[V_1:%.*]] to <8 x i64>
-; CHECK-NEXT:    [[EXT_2:%.*]] = zext <8 x i16> [[V_2:%.*]] to <8 x i64>
-; CHECK-NEXT:    [[SUB:%.*]] = sub nsw <8 x i64> [[EXT_1]], [[EXT_2]]
-; CHECK-NEXT:    [[T:%.*]] = trunc <8 x i64> [[SUB]] to <8 x i32>
-; CHECK-NEXT:    ret <8 x i32> [[T]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <8 x i16> [[V_1:%.*]] to <8 x i32>
+; CHECK-NEXT:    [[TMP2:%.*]] = zext <8 x i16> [[V_2:%.*]] to <8 x i32>
+; CHECK-NEXT:    [[TMP3:%.*]] = sub nsw <8 x i32> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    ret <8 x i32> [[TMP3]]
 ;
   %ext.1 = zext <8 x i16> %v.1 to <8 x i64>
   %ext.2 = zext <8 x i16> %v.2 to <8 x i64>
   %sub = sub nsw <8 x i64> %ext.1, %ext.2
   %t = trunc <8 x i64> %sub to <8 x i32>
   ret <8 x i32> %t
 }
 
 ; Zext can be folded with add for free.
 define <8 x i32> @shorten_zext_add(<8 x i16> %v.1, <8 x i16> %v.2) {
 ; CHECK-LABEL: @shorten_zext_add(
-; CHECK-NEXT:    [[EXT_1:%.*]] = zext <8 x i16> [[V_1:%.*]] to <8 x i64>
-; CHECK-NEXT:    [[EXT_2:%.*]] = zext <8 x i16> [[V_2:%.*]] to <8 x i64>
-; CHECK-NEXT:    [[SUB:%.*]] = add nsw <8 x i64> [[EXT_1]], [[EXT_2]]
-; CHECK-NEXT:    [[T:%.*]] = trunc <8 x i64> [[SUB]] to <8 x i32>
-; CHECK-NEXT:    ret <8 x i32> [[T]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <8 x i16> [[V_1:%.*]] to <8 x i32>
+; CHECK-NEXT:    [[TMP2:%.*]] = zext <8 x i16> [[V_2:%.*]] to <8 x i32>
+; CHECK-NEXT:    [[TMP3:%.*]] = add nsw <8 x i32> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    ret <8 x i32> [[TMP3]]
 ;
   %ext.1 = zext <8 x i16> %v.1 to <8 x i64>
   %ext.2 = zext <8 x i16> %v.2 to <8 x i64>
   %sub = add nsw <8 x i64> %ext.1, %ext.2
   %t = trunc <8 x i64> %sub to <8 x i32>
   ret <8 x i32> %t
 }