[AArch64] Extending lowering of 'trunc <(8|16) x i64> %x to <(8|16) x i8>' to use tbl instructions

[nilanjana_basu]

Summary

[AArch64] Patch for lowering trunc instructions to 'tbl' for (8|16)xi32 -> (8|16)xi8 conversions in D133495 is extended to support trunc to tbl lowering for (8|16) x i64 to (8|16) x i8.

A microbenchmark for runtime has been added for all these cases in D136274

[nilanjana_basu]

Ran git-clang-format & made minor change to reduce LoC.

[nilanjana_basu]

Removed an unused variable warning

[nilanjana_basu]

Extended the trunc lowering for other types like 16xi64, 16xi16, 8xi16

[nilanjana_basu]

The automated build tests failed for the previous patch because it was based on a previous commit for a unit test that isn't submitted yet. This patch fixes it by squashing the previous commit, removing the dependency & showing the final update.

[nilanjana_basu]

Ran clang-format since it was failing in the build report at https://buildkite.com/llvm-project/diff-checks/builds/133184

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13405–13406

I think these are guaranteed to succeed by checks in the caller (and essential here), so cast<...> is probably better. Applies to some of the later dyn_casts too.

13453–13463

There's a lot of duplication in this switch, but it is pretty easy to eyeball for correctness because of that once you get what it's trying to do. So I'm torn, a loop like this would probably be shorter overall:

int ShuffleCount = 128/SrcElemSize;
SmallVector<int> ShuffleLanes;
for (int i = 0; i < ShuffleCount; ++i)
  ShuffleLanes.push_back(i);

SmallVector<Value *> Results;
while (ShuffleLanes.back() < NumElements) {
  Parts.push_back(Builder.CreateBitCast(Builder.CreateShuffleVector(TI->getOperand(0), ShuffleLanes), VecTy));
  for (int i = 0; i < ShuffleCount; ++i)
    ShuffleLanes[i] += ShuffleCount;
  if (Parts.size() == 4) {
    // Call tbl4, push result into Results, clear Parts.
  }
}

// Choose correct tbl (3 now a valid option) and call for rest of Parts, push to Results

// Shuffle-merge all of Results.

and allow the code to apply to a wider range of truncates. What are your views on the implementation?

[nilanjana_basu]

Addressed comments by t.p.northover - refactored code to remove redundancy

[nilanjana_basu]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13453–13463

I refactored the code as you suggested, which can now apply to a few extra cases like 12xi32 or 4xi32. However, I haven't modified the old set of allowable cases since I don't know how relevant these few are.

In my understanding, we get better performance when tbl2-tbl4 get triggered, as the number of generated instructions decrease. So, I need your opinion on whether we should allow 8xi16 conversions, since they generate a single tbl1 instruction?

[nilanjana_basu]

Ran clang-format

[nilanjana_basu]

Added comments

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13405

Is this guaranteed to be a fixed vector type? Could you add a variant of a test with truncates of scalable vectors (<vscale x 16 x i8> or something like that?

13451–13452

It would be great if you could add a brief comment here explaining what kind of masks/shuffles are prepared here.

13475

store here seems ambiguous here, as we won't emit a store instruction, right?

13495

SmallVector?

13500

SmallVector?

llvm/test/CodeGen/AArch64/aarch64-matrix-umull-smull.ll

676 ↗

(On Diff #477022)

Similar to D136722, it is likely not profitable to do this when converting to/from the next power-of-2.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13405

I think it should be fine, I added a test in 4783345426da

[nilanjana_basu]

Updated comments as mentioned in the reviews. Rebased on tests for this change prior to applying this patch.

[nilanjana_basu]

Rebasing on commit of test cases prior to application of this patch

[nilanjana_basu]

Removed case for 'trunc <(8|16)xi16> %x to <(8|16)xi8>' since it was adding more instructions to loop header, while not improving loop instruction count

[nilanjana_basu]

Updated a comment

[nilanjana_basu]

Removed (8|16)xi16 to (8|16)xi8 conversion because it wasn't showing benefits in instruction count, & additionally adding more instructions to the header. Updated comments.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13405

Since the source & destination types were checked for FixedVector once in the calling function optimizeExtendOrTruncateConversion(), I didn't check it here again.

13405

Since both the zext & trunc test for scalable vector goes through the optimizeExtendOrTruncateConversion() function, the zext test in 4783345426da should suffice for the trunc too. Let me know if you think it needs to be replicated for trunc vector too.

13475

I replaced the "store" with "save" to indicate that it is being stored in the compiler's internal vector data structure. Added a comment at the place of combining these results.

llvm/test/CodeGen/AArch64/aarch64-matrix-umull-smull.ll

676 ↗

(On Diff #477022)

Since we only support a handful of vector type truncates in this implementation, only Yxi16->Yxi8 was a supported next power-of-2 conversion. Removed it.

[nilanjana_basu]

Rebasing on parent patch for tests

[nilanjana_basu]

Trying to fix rebasing error

[fhahn]

LGTM with the inline suggestions. Please wait a day or so with committing in case there are additional comments.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13414

could you add an assert to make sure the division happens without remainder?

13436

Could use Builder.getInt8(....)?

13443

IIUC the only case that can happen here is that Parts == 4, right? Might be good to update the check.

[nilanjana_basu]

Addressed the final comments in a separate commit.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static void createTblShuffleForZExt(ZExtInst *ZExt, bool IsLittleEndian) {
   Result = Builder.CreateBitCast(Result, DstTy);
   ZExt->replaceAllUsesWith(Result);
   ZExt->eraseFromParent();
 }
 
 static void createTblForTrunc(TruncInst *TI, bool IsLittleEndian) {
   IRBuilder<> Builder(TI);
   SmallVector<Value *> Parts;
+  unsigned NumElements = cast<FixedVectorType>(TI->getType())->getNumElements();
+  auto *SrcTy = dyn_cast<FixedVectorType>(TI->getOperand(0)->getType());
+  auto *DstTy = dyn_cast<FixedVectorType>(TI->getType());
+  assert(SrcTy->getElementType()->isIntegerTy() &&
+         "Non-integer type source vector element is not supported");
+  assert(DstTy->getElementType()->isIntegerTy(8) &&
+         "Unsupported destination vector element type");
+  unsigned SrcElemTySz =
+      dyn_cast<IntegerType>(SrcTy->getElementType())->getBitWidth();
+  unsigned TruncFactor =
+      SrcElemTySz /
+      dyn_cast<IntegerType>(DstTy->getElementType())->getBitWidth();
+  assert((SrcElemTySz == 32 || SrcElemTySz == 64) &&
+         "Unsupported source vector element type size");
   Type *VecTy = FixedVectorType::get(Builder.getInt8Ty(), 16);
+  Intrinsic::ID TblID = Intrinsic::aarch64_neon_tbl2;
+
+  if (SrcElemTySz == 64 || (SrcElemTySz == 32 && NumElements == 16))
+    TblID = Intrinsic::aarch64_neon_tbl4;
+
+  switch (SrcElemTySz) {
+  case 32:
     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) {
+      TblID = Intrinsic::aarch64_neon_tbl4;
       Parts.push_back(Builder.CreateBitCast(
-        Builder.CreateShuffleVector(TI->getOperand(0), {8, 9, 10, 11}), VecTy));
+          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));
+    }
+    break;
+  case 64:
     TblID = Intrinsic::aarch64_neon_tbl4;
+    Parts.push_back(Builder.CreateBitCast(
+        Builder.CreateShuffleVector(TI->getOperand(0), {0, 1}), VecTy));
+    Parts.push_back(Builder.CreateBitCast(
+        Builder.CreateShuffleVector(TI->getOperand(0), {2, 3}), VecTy));
+    Parts.push_back(Builder.CreateBitCast(
+        Builder.CreateShuffleVector(TI->getOperand(0), {4, 5}), VecTy));
+    Parts.push_back(Builder.CreateBitCast(
+        Builder.CreateShuffleVector(TI->getOperand(0), {6, 7}), VecTy));
+    break;
   }
-  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)
+  SmallVector<Constant *, 16> MaskConst;
+  unsigned Idx = 0;
+  for (unsigned Itr = 0; Itr < 16; Itr++) {
+    if (Itr < NumElements)
+      MaskConst.push_back(ConstantInt::get(
+          Builder.getInt8Ty(), IsLittleEndian
+                                   ? Itr * TruncFactor
+                                   : Itr * TruncFactor + (TruncFactor - 1)));
+    else
       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() ||
       F->hasOptSize())
     return false;
 
   auto *SrcTy = dyn_cast<FixedVectorType>(I->getOperand(0)->getType());
   auto *DstTy = dyn_cast<FixedVectorType>(I->getType());
   if (!SrcTy || !DstTy)
     return false;
 
   // Convert 'zext <(8|16) x i8> %x to <(8|16) x i32>' to a shuffle that can be
   // lowered to either 2 or 4 tbl instructions to insert the original i8
   // elements into i32 lanes.
   auto *ZExt = dyn_cast<ZExtInst>(I);
   if (ZExt && (SrcTy->getNumElements() == 8 || SrcTy->getNumElements() == 16) &&
       SrcTy->getElementType()->isIntegerTy(8) &&
       DstTy->getElementType()->isIntegerTy(32)) {
     createTblShuffleForZExt(ZExt, Subtarget->isLittleEndian());
     return true;
   }
 
   auto *UIToFP = dyn_cast<UIToFPInst>(I);
   if (UIToFP &&
       (SrcTy->getNumElements() == 8 || SrcTy->getNumElements() == 16) &&
       SrcTy->getElementType()->isIntegerTy(8) &&
       DstTy->getElementType()->isFloatTy()) {
     IRBuilder<> Builder(I);
auto *WideConv = Builder.CreateFPToUI(FPToUI->getOperand(0),
                                           VectorType::getInteger(SrcTy));
     auto *TruncI = Builder.CreateTrunc(WideConv, DstTy);
     I->replaceAllUsesWith(TruncI);
     I->eraseFromParent();
     createTblForTrunc(cast<TruncInst>(TruncI), Subtarget->isLittleEndian());
     return true;
   }
 
-  // Convert 'trunc <(8|16) x i32> %x to <(8|16) x i8>' to a single tbl.4
+  // Convert 'trunc <(8|16) x i32> %x to <(8|16) x i8>'
+  // or 'trunc <8 x i64> %x to <8 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)) {
+
+  if (TI && DstTy->getElementType()->isIntegerTy(8) &&
+      ((SrcTy->getElementType()->isIntegerTy(32) &&
+        (SrcTy->getNumElements() == 8 || SrcTy->getNumElements() == 16)) ||
+       (SrcTy->getElementType()->isIntegerTy(64) &&
+        SrcTy->getNumElements() == 8))) {
     createTblForTrunc(TI, Subtarget->isLittleEndian());
     return true;
   }
 
   return false;
 }
 
 bool AArch64TargetLowering::hasPairedLoad(EVT LoadedType,

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

loop:
   %iv.next = add i64 %iv, 1
   %ec = icmp eq i64 %iv.next, 1000
   br i1 %ec, label %loop, label %exit
 
 exit:
   ret void
 }
 
+; CHECK-LABEL: lCPI4_0:
+; CHECK-NEXT:   .byte 0                               ; 0x0
+; CHECK-NEXT:   .byte 8                               ; 0x8
+; CHECK-NEXT:   .byte 16                              ; 0x10
+; CHECK-NEXT:   .byte 24                              ; 0x18
+; CHECK-NEXT:   .byte 32                              ; 0x20
+; CHECK-NEXT:   .byte 40                              ; 0x28
+; CHECK-NEXT:   .byte 48                              ; 0x30
+; CHECK-NEXT:   .byte 56                              ; 0x38
+; 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:   .LCPI4_0:
+; CHECK-BE-NEXT:    .byte 7                               // 0x7
+; CHECK-BE-NEXT:    .byte 15                              // 0xf
+; CHECK-BE-NEXT:    .byte 23                              // 0x17
+; CHECK-BE-NEXT:    .byte 31                              // 0x1f
+; CHECK-BE-NEXT:    .byte 39                              // 0x27
+; CHECK-BE-NEXT:    .byte 47                              // 0x2f
+; CHECK-BE-NEXT:    .byte 55                              // 0x37
+; CHECK-BE-NEXT:    .byte 63                              // 0x3f
+; 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
 define void @trunc_v8i64_to_v8i8_in_loop(ptr %A, ptr %dst) {
 ; CHECK-LABEL:  trunc_v8i64_to_v8i8_in_loop:
-; CHECK:    ; %bb.0:                                ; %entry
-; CHECK-NEXT:   mov x8, xzr
-; CHECK-NEXT: LBB4_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.4s v0, v1, v0
-; CHECK-NEXT:   uzp1.4s v1, v3, v2
-; CHECK-NEXT:   uzp1.8h v0, v1, v0
-; CHECK-NEXT:   xtn.8b  v0, v0
-; CHECK-NEXT:   str d0, [x1, x8, lsl  #3]
-; CHECK-NEXT:   add x8, x8, #1
-; CHECK-NEXT:   cmp x8, #1000
-; CHECK-NEXT:   b.eq  LBB4_1
-; CHECK-NEXT:  ; %bb.2:                                ; %exit
-; CHECK-NEXT:   ret
+; CHECK:    ; %bb.0:                                ; %entry
+; CHECK-NEXT:   Lloh4:
+; CHECK-NEXT:     adrp  x9, lCPI4_0@PAGE
+; CHECK-NEXT:     mov x8, xzr
+; CHECK-NEXT:   Lloh5:
+; CHECK-NEXT:     ldr q0, [x9, lCPI4_0@PAGEOFF]
+; CHECK-NEXT:   LBB4_1:                                 ; %loop
+; CHECK-NEXT:                                           ; =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:     add x9, x0, x8, lsl #6
+; 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 d1, [x1, x8, lsl  #3]
+; CHECK-NEXT:     add x8, x8, #1
+; CHECK-NEXT:     cmp x8, #1000
+; CHECK-NEXT:     b.eq  LBB4_1
+; CHECK-NEXT:   ; %bb.2:                                ; %exit
+; CHECK-NEXT:     ret
+; CHECK-NEXT:     .loh AdrpLdr  Lloh4, Lloh5
 
 ; CHECK-BE-LABEL:  trunc_v8i64_to_v8i8_in_loop:
 ; CHECK-BE:  // %bb.0:                               // %entry
-; CHECK-BE-NEXT:    mov x8, xzr
-; CHECK-BE-NEXT: .LBB4_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:   ld1  { v1.2d }, [x9]
-; CHECK-BE-NEXT:   ld1  { v0.2d }, [x10]
-; CHECK-BE-NEXT:   add  x10, x9, #32
-; CHECK-BE-NEXT:   add  x9, x9, #16
-; CHECK-BE-NEXT:   ld1  { v2.2d }, [x10]
-; CHECK-BE-NEXT:   ld1  { v3.2d }, [x9]
-; CHECK-BE-NEXT:   add  x9, x1, x8, lsl #3
-; CHECK-BE-NEXT:   add  x8, x8, #1
-; CHECK-BE-NEXT:   cmp  x8, #1000
-; CHECK-BE-NEXT:   uzp1 v0.4s, v2.4s, v0.4s
-; CHECK-BE-NEXT:   uzp1 v1.4s, v1.4s, v3.4s
-; CHECK-BE-NEXT:   uzp1 v0.8h, v1.8h, v0.8h
-; CHECK-BE-NEXT:   xtn  v0.8b, v0.8h
-; CHECK-BE-NEXT:   st1  { v0.8b }, [x9]
-; CHECK-BE-NEXT:   b.eq .LBB4_1
+; CHECK-BE-NEXT:    adrp  x8, .LCPI4_0
+; CHECK-BE-NEXT:    add x8, x8, :lo12:.LCPI4_0
+; CHECK-BE-NEXT:    ld1 { v0.16b }, [x8]
+; CHECK-BE-NEXT:    mov x8, xzr
+; CHECK-BE-NEXT:  .LBB4_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, #16
+; CHECK-BE-NEXT:    add x11, x9, #32
+; 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 #3
+; CHECK-BE-NEXT:    add x8, x8, #1
+; CHECK-BE-NEXT:    cmp x8, #1000
+; CHECK-BE-NEXT:    tbl v1.16b, { v1.16b, v2.16b, v3.16b, v4.16b }, v0.16b
+; CHECK-BE-NEXT:    st1 { v1.8b }, [x9]
+; CHECK-BE-NEXT:    b.eq  .LBB4_1
 ; CHECK-BE-NEXT:   // %bb.2:                               // %exit
-; CHECK-BE-NEXT:    ret
+; CHECK-BE-NEXT:    ret
 
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
   %gep.A = getelementptr inbounds <8 x i64>, ptr %A, i64 %iv
   %l.A = load <8 x i64>, ptr %gep.A