[ARM] Allow tail predication on sadd_sat and uadd_sat intrinsics

[samtebbs]

Summary

This patch stops the sadd_sat and uadd_sat intrinsics from blocking tail predication.

[SjoerdMeijer]

llvm/lib/Target/ARM/MVETailPredication.cpp

355–356

Do we expect more intrinsics to be allowed after these 2? For example, if we allow sadd, can we allow ssub too? And then while we are add it, add the ssubs too?
If we expect more intrinsics here, probably better to create an "AllowList" for intrinsics somewhere, and do a look-up here?

[dmgreen]

llvm/lib/Target/ARM/MVETailPredication.cpp

355–356

Can you give this a clang-format? The line is a bit long.

We might end up with quite a few of these too. It might make sense to add a switch on Int->getIntrinsicID() once this gets large enough.

llvm/test/CodeGen/Thumb2/LowOverheadLoops/tail-pred-intrinsic-uadd.ll

3 ↗

(On Diff #272696)

You can remove dso_local and local_unnamed_addr.

I would also make the two tests one file to keep them together, but up to you if you want to keep them separate.

[xbolva00]

llvm/lib/Target/ARM/MVETailPredication.cpp

355–356

+1 for switch

[samtebbs]

llvm/lib/Target/ARM/MVETailPredication.cpp

355–356

I could replace this with a switch statement as suggested. What do you think?

355–356

Agreed. A switch would be much better.

llvm/test/CodeGen/Thumb2/LowOverheadLoops/tail-pred-intrinsic-uadd.ll

3 ↗

(On Diff #272696)

I'll remove ds_local and local_unnamed and merge them +1

[samtebbs]

Change to switch statement and merge tests

[samtebbs]

Remove dso_local and local_unnamed from test

[dmgreen]

Thanks. LGTM

[samtebbs]

Closed by 187f627a505

llvm/lib/Target/ARM/MVETailPredication.cpp

for (auto &I : *BB) {
         unsigned Lanes = VecTy->getNumElements();
         unsigned ElementWidth = VecTy->getScalarSizeInBits();
         // MVE vectors are 128-bit, but don't support 128 x i1.
         // TODO: Can we support vectors larger than 128-bits?
         unsigned MaxWidth = TTI->getRegisterBitWidth(true);
         if (Lanes * ElementWidth > MaxWidth || Lanes == MaxWidth)
           return false;
         MaskedInsts.push_back(cast<IntrinsicInst>(&I));
       } else if (auto *Int = dyn_cast<IntrinsicInst>(&I)) {
-        if (Int->getIntrinsicID() == Intrinsic::fma)
+        switch (Int->getIntrinsicID()) {
+          case Intrinsic::fma:
+          case Intrinsic::sadd_sat:
+          case Intrinsic::uadd_sat:
             continue;
+          default:
+            break;
+        }
         for (auto &U : Int->args()) {
           if (isa<VectorType>(U->getType()))
             return false;
         }
       }
     }
   }
 

llvm/test/CodeGen/Thumb2/LowOverheadLoops/tail-pred-intrinsic-add-sat.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+mve -verify-machineinstrs -disable-mve-tail-predication=false -o - %s | FileCheck %s
+define dso_local arm_aapcs_vfpcc void @uadd_sat(i16* noalias nocapture readonly %pSrcA, i16* noalias nocapture readonly %pSrcB, i16* noalias nocapture %pDst, i32 %blockSize) local_unnamed_addr {
+; CHECK-LABEL: uadd_sat:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    .save {r7, lr}
+; CHECK-NEXT:    push {r7, lr}
+; CHECK-NEXT:    cmp r3, #0
+; CHECK-NEXT:    it eq
+; CHECK-NEXT:    popeq {r7, pc}
+; CHECK-NEXT:    dlstp.16 lr, r3
+; CHECK-NEXT:  .LBB0_1: @ %vector.body
+; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    vldrh.u16 q0, [r1], #16
+; CHECK-NEXT:    vldrh.u16 q1, [r0], #16
+; CHECK-NEXT:    vqadd.u16 q0, q1, q0
+; CHECK-NEXT:    vstrh.16 q0, [r2], #16
+; CHECK-NEXT:    letp lr, .LBB0_1
+; CHECK-NEXT:  @ %bb.2: @ %while.end
+; CHECK-NEXT:    pop {r7, pc}
+entry:
+  %cmp7 = icmp eq i32 %blockSize, 0
+  br i1 %cmp7, label %while.end, label %vector.ph
+
+vector.ph:                                        ; preds = %entry
+  %n.rnd.up = add i32 %blockSize, 7
+  %n.vec = and i32 %n.rnd.up, -8
+  %trip.count.minus.1 = add i32 %blockSize, -1
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %next.gep = getelementptr i16, i16* %pSrcA, i32 %index
+  %next.gep20 = getelementptr i16, i16* %pDst, i32 %index
+  %next.gep21 = getelementptr i16, i16* %pSrcB, i32 %index
+  %active.lane.mask = call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 %index, i32 %trip.count.minus.1)
+  %0 = bitcast i16* %next.gep to <8 x i16>*
+  %wide.masked.load = call <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>* %0, i32 2, <8 x i1> %active.lane.mask, <8 x i16> undef)
+  %1 = bitcast i16* %next.gep21 to <8 x i16>*
+  %wide.masked.load24 = call <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>* %1, i32 2, <8 x i1> %active.lane.mask, <8 x i16> undef)
+  %2 = call <8 x i16> @llvm.uadd.sat.v8i16(<8 x i16> %wide.masked.load, <8 x i16> %wide.masked.load24)
+  %3 = bitcast i16* %next.gep20 to <8 x i16>*
+  call void @llvm.masked.store.v8i16.p0v8i16(<8 x i16> %2, <8 x i16>* %3, i32 2, <8 x i1> %active.lane.mask)
+  %index.next = add i32 %index, 8
+  %4 = icmp eq i32 %index.next, %n.vec
+  br i1 %4, label %while.end, label %vector.body
+
+while.end:                                        ; preds = %vector.body, %entry
+  ret void
+}
+
+define dso_local arm_aapcs_vfpcc void @sadd_sat(i16* noalias nocapture readonly %pSrcA, i16* noalias nocapture readonly %pSrcB, i16* noalias nocapture %pDst, i32 %blockSize) local_unnamed_addr {
+; CHECK-LABEL: sadd_sat:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    .save {r7, lr}
+; CHECK-NEXT:    push {r7, lr}
+; CHECK-NEXT:    cmp r3, #0
+; CHECK-NEXT:    it eq
+; CHECK-NEXT:    popeq {r7, pc}
+; CHECK-NEXT:    dlstp.16 lr, r3
+; CHECK-NEXT:  .LBB1_1: @ %vector.body
+; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    vldrh.u16 q0, [r1], #16
+; CHECK-NEXT:    vldrh.u16 q1, [r0], #16
+; CHECK-NEXT:    vqadd.s16 q0, q1, q0
+; CHECK-NEXT:    vstrh.16 q0, [r2], #16
+; CHECK-NEXT:    letp lr, .LBB1_1
+; CHECK-NEXT:  @ %bb.2: @ %while.end
+; CHECK-NEXT:    pop {r7, pc}
+entry:
+  %cmp7 = icmp eq i32 %blockSize, 0
+  br i1 %cmp7, label %while.end, label %vector.ph
+
+vector.ph:                                        ; preds = %entry
+  %n.rnd.up = add i32 %blockSize, 7
+  %n.vec = and i32 %n.rnd.up, -8
+  %trip.count.minus.1 = add i32 %blockSize, -1
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %next.gep = getelementptr i16, i16* %pSrcA, i32 %index
+  %next.gep20 = getelementptr i16, i16* %pDst, i32 %index
+  %next.gep21 = getelementptr i16, i16* %pSrcB, i32 %index
+  %active.lane.mask = call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 %index, i32 %trip.count.minus.1)
+  %0 = bitcast i16* %next.gep to <8 x i16>*
+  %wide.masked.load = call <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>* %0, i32 2, <8 x i1> %active.lane.mask, <8 x i16> undef)
+  %1 = bitcast i16* %next.gep21 to <8 x i16>*
+  %wide.masked.load24 = call <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>* %1, i32 2, <8 x i1> %active.lane.mask, <8 x i16> undef)
+  %2 = call <8 x i16> @llvm.sadd.sat.v8i16(<8 x i16> %wide.masked.load, <8 x i16> %wide.masked.load24)
+  %3 = bitcast i16* %next.gep20 to <8 x i16>*
+  call void @llvm.masked.store.v8i16.p0v8i16(<8 x i16> %2, <8 x i16>* %3, i32 2, <8 x i1> %active.lane.mask)
+  %index.next = add i32 %index, 8
+  %4 = icmp eq i32 %index.next, %n.vec
+  br i1 %4, label %while.end, label %vector.body
+
+while.end:                                        ; preds = %vector.body, %entry
+  ret void
+}
+
+declare <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32, i32)
+
+declare <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>*, i32 immarg, <8 x i1>, <8 x i16>)
+
+declare <8 x i16> @llvm.sadd.sat.v8i16(<8 x i16>, <8 x i16>)
+
+declare <8 x i16> @llvm.uadd.sat.v8i16(<8 x i16>, <8 x i16>)
+
+declare void @llvm.masked.store.v8i16.p0v8i16(<8 x i16>, <8 x i16>*, i32 immarg, <8 x i1>)