[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

13216–13217

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.

13289

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

13292

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

13216–13217

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

13289

Yes, moved to D133494.

13292

Thanks, should be fixed!

[t.p.northover]

Thanks Florian, think this one looks good now.

llvm/lib/CodeGen/CodeGenPrepare.cpp

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) {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static void createTblShuffleForZExt(ZExtInst *ZExt, bool IsLittleEndian) {
   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, bool IsLittleEndian) {
+  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(), IsLittleEndian ? Idx : Idx + 3));
+
+  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() ||
if (UIToFP &&
       DstTy->getElementType()->isFloatTy()) {
     IRBuilder<> Builder(I);
     auto *ZExt = cast<ZExtInst>(
         Builder.CreateZExt(I->getOperand(0), VectorType::getInteger(DstTy)));
     auto *UI = Builder.CreateUIToFP(ZExt, DstTy);
     I->replaceAllUsesWith(UI);
     I->eraseFromParent();
     createTblShuffleForZExt(ZExt, Subtarget->isLittleEndian());
     return true;
   }
+
+  // Convert 'trunc <(8|16) x i32> %x to <(8|16) x i8>' 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, Subtarget->isLittleEndian());
+    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 -mtriple=arm64-apple-ios -o - %s | FileCheck %s
 ; RUN: llc -mtriple=aarch64_be-unknown-linux -o - %s | FileCheck --check-prefix=CHECK-BE %s
 
+; 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
+
+; CHECK-BE-LABEL:   .LCPI0_0:
+; CHECK-BE-NEXT:   .byte    3                               // 0x3
+; CHECK-BE-NEXT:   .byte    7                               // 0x7
+; CHECK-BE-NEXT:   .byte    11                              // 0xb
+; CHECK-BE-NEXT:   .byte    15                              // 0xf
+; CHECK-BE-NEXT:   .byte    19                              // 0x13
+; CHECK-BE-NEXT:   .byte    23                              // 0x17
+; CHECK-BE-NEXT:   .byte    27                              // 0x1b
+; CHECK-BE-NEXT:   .byte    31                              // 0x1f
+; CHECK-BE-NEXT:   .byte    35                              // 0x23
+; CHECK-BE-NEXT:   .byte    39                              // 0x27
+; CHECK-BE-NEXT:   .byte    43                              // 0x2b
+; CHECK-BE-NEXT:   .byte    47                              // 0x2f
+; CHECK-BE-NEXT:   .byte    51                              // 0x33
+; CHECK-BE-NEXT:   .byte    55                              // 0x37
+; CHECK-BE-NEXT:   .byte    59                              // 0x3b
+; CHECK-BE-NEXT:   .byte    63                              // 0x3f
+
 ; 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
 ;
 ; CHECK-BE-LABEL: trunc_v16i32_to_v16i8_in_loop:
 ; CHECK-BE:       // %bb.0: // %entry
+; CHECK-BE-NEXT:    adrp x8, .LCPI0_0
+; CHECK-BE-NEXT:    add x8, x8, :lo12:.LCPI0_0
+; CHECK-BE-NEXT:    ld1 { v0.16b }, [x8]
 ; CHECK-BE-NEXT:    mov x8, xzr
 ; CHECK-BE-NEXT:  .LBB0_1: // %loop
 ; CHECK-BE-NEXT:    // =>This Inner Loop Header: Depth=1
 ; CHECK-BE-NEXT:    add x9, x0, x8, lsl #6
-; CHECK-BE-NEXT:    add x10, x9, #48
+; CHECK-BE-NEXT:    add x10, x9, #16
 ; CHECK-BE-NEXT:    add x11, x9, #32
-; CHECK-BE-NEXT:    ld1 { v0.4s }, [x9]
-; CHECK-BE-NEXT:    add x9, x9, #16
-; CHECK-BE-NEXT:    ld1 { v1.4s }, [x10]
-; CHECK-BE-NEXT:    ld1 { v2.4s }, [x11]
-; CHECK-BE-NEXT:    ld1 { v3.4s }, [x9]
+; CHECK-BE-NEXT:    ld1 { v1.16b }, [x9]
+; CHECK-BE-NEXT:    add x9, x9, #48
+; CHECK-BE-NEXT:    ld1 { v2.16b }, [x10]
+; CHECK-BE-NEXT:    ld1 { v3.16b }, [x11]
+; CHECK-BE-NEXT:    ld1 { v4.16b }, [x9]
 ; CHECK-BE-NEXT:    add x9, x1, x8, lsl #4
 ; CHECK-BE-NEXT:    add x8, x8, #1
 ; CHECK-BE-NEXT:    cmp x8, #1000
-; CHECK-BE-NEXT:    uzp1 v1.8h, v2.8h, v1.8h
-; CHECK-BE-NEXT:    uzp1 v0.8h, v0.8h, v3.8h
-; CHECK-BE-NEXT:    uzp1 v0.16b, v0.16b, v1.16b
-; CHECK-BE-NEXT:    st1 { v0.16b }, [x9]
+; CHECK-BE-NEXT:    tbl v1.16b, { v1.16b, v2.16b, v3.16b, v4.16b }, v0.16b
+; CHECK-BE-NEXT:    st1 { v1.16b }, [x9]
 ; CHECK-BE-NEXT:    b.eq .LBB0_1
 ; CHECK-BE-NEXT:  // %bb.2: // %exit
 ; CHECK-BE-NEXT:    ret
+
 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-BE-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
+
+; CHECK-BE-LABEL: .LCPI2_0:
+; CHECK-BE-NEXT:     .byte    3                               // 0x3
+; CHECK-BE-NEXT:     .byte    7                               // 0x7
+; CHECK-BE-NEXT:     .byte    11                              // 0xb
+; CHECK-BE-NEXT:     .byte    15                              // 0xf
+; CHECK-BE-NEXT:     .byte    19                              // 0x13
+; CHECK-BE-NEXT:     .byte    23                              // 0x17
+; CHECK-BE-NEXT:     .byte    27                              // 0x1b
+; CHECK-BE-NEXT:     .byte    31                              // 0x1f
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-NEXT:     .byte    255                             // 0xff
+; CHECK-BE-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
 ;
 ; CHECK-BE-LABEL: trunc_v8i32_to_v8i8_in_loop:
 ; CHECK-BE:       // %bb.0: // %entry
+; CHECK-BE-NEXT:    adrp x8, .LCPI2_0
+; CHECK-BE-NEXT:    add x8, x8, :lo12:.LCPI2_0
+; CHECK-BE-NEXT:    ld1 { v0.16b }, [x8]
 ; CHECK-BE-NEXT:    mov x8, xzr
 ; CHECK-BE-NEXT:  .LBB2_1: // %loop
 ; CHECK-BE-NEXT:    // =>This Inner Loop Header: Depth=1
 ; CHECK-BE-NEXT:    add x9, x0, x8, lsl #5
 ; CHECK-BE-NEXT:    add x10, x9, #16
-; CHECK-BE-NEXT:    ld1 { v0.4s }, [x9]
+; CHECK-BE-NEXT:    ld1 { v1.16b }, [x9]
 ; CHECK-BE-NEXT:    add x9, x1, x8, lsl #3
 ; CHECK-BE-NEXT:    add x8, x8, #1
-; CHECK-BE-NEXT:    ld1 { v1.4s }, [x10]
+; CHECK-BE-NEXT:    ld1 { v2.16b }, [x10]
 ; CHECK-BE-NEXT:    cmp x8, #1000
-; CHECK-BE-NEXT:    uzp1 v0.8h, v0.8h, v1.8h
-; CHECK-BE-NEXT:    xtn v0.8b, v0.8h
-; CHECK-BE-NEXT:    st1 { v0.8b }, [x9]
+; CHECK-BE-NEXT:    tbl v1.16b, { v1.16b, v2.16b }, v0.16b
+; CHECK-BE-NEXT:    st1 { v1.8b }, [x9]
 ; CHECK-BE-NEXT:    b.eq .LBB2_1
 ; CHECK-BE-NEXT:  // %bb.2: // %exit
 ; CHECK-BE-NEXT:    ret
+
 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>