[AArch64] Lower vector trunc using tbl.

[fhahn]

Summary

Similar to using tbl to lower vector ZExts, tbl4 can be used to lower
vector truncates.

The initial version support i32->i8 conversions.

Depends on D120571

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13246–13247

Further to our big-endian discussion, I think this bitcast will put the least significant byte into elements 3, 7, 11, 15. So the for loop below should account for that.

13324

Shouldn't this be part of the patch that adds the code above?

13327

Destination type wrong.

[fhahn]

Thanks Tim, I updated the code to account for big-endian targets by adjusting the mask as suggested offline.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13246–13247

Thanks, I adjusted the loop below to use Idx+3 for the mask to account for that on big-endian.

13324

Yes, moved to D133494.

13327

Thanks, should be fixed!

[t.p.northover]

Thanks Florian, think this one looks good now.

llvm/lib/CodeGen/CodeGenPrepare.cpp

for (auto *VisitedSExt : UnhandledExts) {
         ValToSExtendedUses[HeadOfChain].push_back(I);
       }
     }
   return Promoted;
 }
 
 bool CodeGenPrepare::optimizeExtUses(Instruction *I) {
   BasicBlock *DefBB = I->getParent();
-
   // If the result of a {s|z}ext and its source are both live out, rewrite all
   // other uses of the source with result of extension.
   Value *Src = I->getOperand(0);
   if (Src->hasOneUse())
     return false;
 
   // Only do this xform if truncating is free.
   if (!TLI->isTruncateFree(I->getType(), Src->getType()))
if (CastInst *CI = dyn_cast<CastInst>(I)) {
     // the address of globals out of a loop).  If this is the case, we don't
     // want to forward-subst the cast.
     if (isa<Constant>(CI->getOperand(0)))
       return false;
 
     if (OptimizeNoopCopyExpression(CI, *TLI, *DL))
       return true;
 
-    if (isa<UIToFPInst>(I) && TLI->optimizeExtendOrTruncateConversion(
-                                  I, LI->getLoopFor(I->getParent())))
+    if ((isa<UIToFPInst>(I) || isa<TruncInst>(I)) &&
+        TLI->optimizeExtendOrTruncateConversion(I,
+                                                LI->getLoopFor(I->getParent())))
       return true;
 
     if (isa<ZExtInst>(I) || isa<SExtInst>(I)) {
       /// Sink a zext or sext into its user blocks if the target type doesn't
       /// fit in one register
       if (TLI->getTypeAction(CI->getContext(),
                              TLI->getValueType(*DL, CI->getType())) ==
           TargetLowering::TypeExpandInteger) {
         return SinkCast(CI);
       } else {
         if (TLI->optimizeExtendOrTruncateConversion(
                 I, LI->getLoopFor(I->getParent())))
           return true;
 
         bool MadeChange = optimizeExt(I);
         return MadeChange | optimizeExtUses(I);
       }
     }
+
     return false;
   }
 
   if (auto *Cmp = dyn_cast<CmpInst>(I))
     if (optimizeCmp(Cmp, ModifiedDT))
       return true;
 
   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static void createTblShuffleForZExt(ZExtInst *ZExt) {
   auto *FirstEltZero = Builder.CreateInsertElement(
       PoisonValue::get(SrcTy), Builder.getInt8(0), uint64_t(0));
   Value *Result = Builder.CreateShuffleVector(Op, FirstEltZero, Mask);
   Result = Builder.CreateBitCast(Result, DstTy);
   ZExt->replaceAllUsesWith(Result);
   ZExt->eraseFromParent();
 }
 
+static void createTblForTrunc(TruncInst *TI) {
+  IRBuilder<> Builder(TI);
+  SmallVector<Value *> Parts;
+  Type *VecTy = FixedVectorType::get(Builder.getInt8Ty(), 16);
+  Parts.push_back(Builder.CreateBitCast(
+      Builder.CreateShuffleVector(TI->getOperand(0), {0, 1, 2, 3}), VecTy));
+  Parts.push_back(Builder.CreateBitCast(
+      Builder.CreateShuffleVector(TI->getOperand(0), {4, 5, 6, 7}), VecTy));
+
+  Intrinsic::ID TblID = Intrinsic::aarch64_neon_tbl2;
+  unsigned NumElements = cast<FixedVectorType>(TI->getType())->getNumElements();
+  if (NumElements == 16) {
+    Parts.push_back(Builder.CreateBitCast(
+        Builder.CreateShuffleVector(TI->getOperand(0), {8, 9, 10, 11}), VecTy));
+    Parts.push_back(Builder.CreateBitCast(
+        Builder.CreateShuffleVector(TI->getOperand(0), {12, 13, 14, 15}),
+        VecTy));
+    TblID = Intrinsic::aarch64_neon_tbl4;
+  }
+  SmallVector<Constant *, 16> MaskConst;
+  for (unsigned Idx = 0; Idx < NumElements * 4; Idx += 4)
+    MaskConst.push_back(ConstantInt::get(Builder.getInt8Ty(), Idx));
+
+  for (unsigned Idx = NumElements * 4; Idx < 64; Idx += 4)
+    MaskConst.push_back(ConstantInt::get(Builder.getInt8Ty(), 255));
+
+  Parts.push_back(ConstantVector::get(MaskConst));
+  auto *F =
+      Intrinsic::getDeclaration(TI->getModule(), TblID, Parts[0]->getType());
+  Value *Res = Builder.CreateCall(F, Parts);
+
+  if (NumElements == 8)
+    Res = Builder.CreateShuffleVector(Res, {0, 1, 2, 3, 4, 5, 6, 7});
+  TI->replaceAllUsesWith(Res);
+  TI->eraseFromParent();
+}
+
 bool AArch64TargetLowering::optimizeExtendOrTruncateConversion(Instruction *I,
                                                                Loop *L) const {
   // Try to optimize conversions using tbl. This requires materializing constant
   // index vectors, which can increase code size and add loads. Skip the
   // transform unless the conversion is in a loop block guaranteed to execute
   // and we are not optimizing for size.
   Function *F = I->getParent()->getParent();
   if (!L || L->getHeader() != I->getParent() || F->hasMinSize() ||
auto *ZExt = cast<ZExtInst>(Builder.CreateZExt(
                 I->getType()->getScalarType()->getScalarSizeInBits()),
             cast<FixedVectorType>(I->getType())
                 ->getElementCount()
                 .getKnownMinValue())));
     auto *UI = Builder.CreateUIToFP(ZExt, I->getType());
     I->replaceAllUsesWith(UI);
     I->eraseFromParent();
     createTblShuffleForZExt(ZExt);
+    return true;
   }
+
+  // Convert 'trunc <(8|16) x i32> %x to <(8|16) x i32>' to a single tbl.4
+  // instruction selecting the lowest 8 bits per lane of the input interpreted
+  // as 2 or 4 <4 x i32> vectors.
+  auto *TI = dyn_cast<TruncInst>(I);
+  if (TI && (SrcTy->getNumElements() == 8 || SrcTy->getNumElements() == 16) &&
+      SrcTy->getElementType()->isIntegerTy(32) &&
+      DstTy->getElementType()->isIntegerTy(8)) {
+    createTblForTrunc(TI);
+    return true;
+  }
+
   return false;
 }
 
 bool AArch64TargetLowering::hasPairedLoad(EVT LoadedType,
                                           Align &RequiredAligment) const {
   if (!LoadedType.isSimple() ||
       (!LoadedType.isInteger() && !LoadedType.isFloatingPoint()))
     return false;

llvm/test/CodeGen/AArch64/trunc-to-tbl.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -o - %s | FileCheck %s
 
 target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
 target triple = "arm64-apple-ios"
 
+; CHECK-LABEL: lCPI0_0:
+; CHECK-NEXT: .byte 0                               ; 0x0
+; CHECK-NEXT: .byte 4                               ; 0x4
+; CHECK-NEXT: .byte 8                               ; 0x8
+; CHECK-NEXT: .byte 12                              ; 0xc
+; CHECK-NEXT: .byte 16                              ; 0x10
+; CHECK-NEXT: .byte 20                              ; 0x14
+; CHECK-NEXT: .byte 24                              ; 0x18
+; CHECK-NEXT: .byte 28                              ; 0x1c
+; CHECK-NEXT: .byte 32                              ; 0x20
+; CHECK-NEXT: .byte 36                              ; 0x24
+; CHECK-NEXT: .byte 40                              ; 0x28
+; CHECK-NEXT: .byte 44                              ; 0x2c
+; CHECK-NEXT: .byte 48                              ; 0x30
+; CHECK-NEXT: .byte 52                              ; 0x34
+; CHECK-NEXT: .byte 56                              ; 0x38
+; CHECK-NEXT: .byte 60                              ; 0x3c
+
+
 ; It's profitable to use a single tbl.4 instruction to lower the truncate.
 define void @trunc_v16i32_to_v16i8_in_loop(ptr %A, ptr %dst) {
 ; CHECK-LABEL: trunc_v16i32_to_v16i8_in_loop:
 ; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:  Lloh0:
+; CHECK-NEXT:    adrp x9, lCPI0_0@PAGE
 ; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:  Lloh1:
+; CHECK-NEXT:    ldr q0, [x9, lCPI0_0@PAGEOFF]
 ; CHECK-NEXT:  LBB0_1: ; %loop
 ; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    add x9, x0, x8, lsl #6
-; CHECK-NEXT:    ldp q1, q0, [x9, #32]
-; CHECK-NEXT:    ldp q3, q2, [x9]
-; CHECK-NEXT:    uzp1.8h v0, v1, v0
-; CHECK-NEXT:    uzp1.8h v1, v3, v2
-; CHECK-NEXT:    uzp1.16b v0, v1, v0
-; CHECK-NEXT:    str q0, [x1, x8, lsl #4]
+; CHECK-NEXT:    ldp q1, q2, [x9]
+; CHECK-NEXT:    ldp q3, q4, [x9, #32]
+; CHECK-NEXT:    tbl.16b v1, { v1, v2, v3, v4 }, v0
+; CHECK-NEXT:    str q1, [x1, x8, lsl #4]
 ; CHECK-NEXT:    add x8, x8, #1
 ; CHECK-NEXT:    cmp x8, #1000
 ; CHECK-NEXT:    b.eq LBB0_1
 ; CHECK-NEXT:  ; %bb.2: ; %exit
 ; CHECK-NEXT:    ret
+; CHECK-NEXT:    .loh AdrpLdr Lloh0, Lloh1
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %gep.A = getelementptr inbounds <16 x i32>, ptr %A, i64 %iv
   %l.A = load <16 x i32>, ptr %gep.A
   %trunc = trunc <16 x i32> %l.A to <16 x i8>

 ; CHECK-NEXT:    ret
 entry:
   %l.A = load <16 x i32>, ptr %A
   %trunc = trunc <16 x i32> %l.A to <16 x i8>
   store <16 x i8> %trunc, ptr %dst
   ret void
 }
 
+
+; CHECK-LABEL: lCPI2_0:
+; CHECK-NEXT:     .byte    0                               ; 0x0
+; CHECK-NEXT:     .byte    4                               ; 0x4
+; CHECK-NEXT:     .byte    8                               ; 0x8
+; CHECK-NEXT:     .byte    12                              ; 0xc
+; CHECK-NEXT:     .byte    16                              ; 0x10
+; CHECK-NEXT:     .byte    20                              ; 0x14
+; CHECK-NEXT:     .byte    24                              ; 0x18
+; CHECK-NEXT:     .byte    28                              ; 0x1c
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+; CHECK-NEXT:     .byte    255                             ; 0xff
+
 ; It's profitable to use a single tbl.2 instruction to lower the truncate.
 define void @trunc_v8i32_to_v8i8_in_loop(ptr %A, ptr %dst) {
 ; CHECK-LABEL: trunc_v8i32_to_v8i8_in_loop:
 ; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:  Lloh2:
+; CHECK-NEXT:    adrp x9, lCPI2_0@PAGE
 ; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:  Lloh3:
+; CHECK-NEXT:    ldr q0, [x9, lCPI2_0@PAGEOFF]
 ; CHECK-NEXT:  LBB2_1: ; %loop
 ; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    add x9, x0, x8, lsl #5
-; CHECK-NEXT:    ldp q1, q0, [x9]
-; CHECK-NEXT:    uzp1.8h v0, v1, v0
-; CHECK-NEXT:    xtn.8b v0, v0
-; CHECK-NEXT:    str d0, [x1, x8, lsl #3]
+; CHECK-NEXT:    ldp q1, q2, [x9]
+; CHECK-NEXT:    tbl.16b v1, { v1, v2 }, v0
+; CHECK-NEXT:    str d1, [x1, x8, lsl #3]
 ; CHECK-NEXT:    add x8, x8, #1
 ; CHECK-NEXT:    cmp x8, #1000
 ; CHECK-NEXT:    b.eq LBB2_1
 ; CHECK-NEXT:  ; %bb.2: ; %exit
 ; CHECK-NEXT:    ret
+; CHECK-NEXT:    .loh AdrpLdr Lloh2, Lloh3
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %gep.A = getelementptr inbounds <8 x i32>, ptr %A, i64 %iv
   %l.A = load <8 x i32>, ptr %gep.A
   %trunc = trunc <8 x i32> %l.A to <8 x i8>