[ARM][MVE] Tail-Predication: recognise (again) active lanes IR pattern

[SjoerdMeijer]

Summary

A small IR change in calculating the active lanes resulted in no longer recognising tail-predication. We used to calculate a vector containing the iteration counts for a vector body with e.g. 4 iterations with:

[i, i, i, i] + [0, 1, 2, 3]

but this is changed to:

[i, i, i, i] or [0, 1, 2, 3]

when this result is compared to the vector with loop iteration count vector [N, N, N, N] then this is equivalent, so now recognise both the 'add' and 'or'.

[dmgreen]

LGTM, if you can add a comment too.

llvm/lib/Target/ARM/MVETailPredication.cpp

273–274

I think if this was using an or it was because the BroadcastSplat value is mod Lanes. That is OK here because BroadcastSplat is IV which we know is a phi with 0 start and Lanes increment?

If so, can you add a comment that we know that it will be aligned, so an Or is equivalent to an Add?

[samparker]

llvm/test/CodeGen/Thumb2/LowOverheadLoops/basic-tail-pred.ll

35

Does this mean we're no longer directly testing the 'add' form?

[SjoerdMeijer]

if you can add a comment too.

thanks, will do

llvm/test/CodeGen/Thumb2/LowOverheadLoops/basic-tail-pred.ll

88

Does this mean we're no longer directly testing the 'add' form?

We are testing both, there's one here for example, and I alternated the test here using 'and' and 'or'

llvm/lib/Target/ARM/MVETailPredication.cpp

bool MVETailPredication::isTailPredicate(TripCountPattern &TCP) {
   //
   // vector.body:
   // ..
   // %index = phi i32
   // %broadcast.splatinsert = insertelement <4 x i32> undef, i32 %index, i32 0
   // %broadcast.splat = shufflevector <4 x i32> %broadcast.splatinsert,
   //                                  <4 x i32> undef,
   //                                  <4 x i32> zeroinitializer
-  // %induction = add <4 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3>
+  // %induction = [add|or] <4 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3>
   // %pred = icmp ule <4 x i32> %induction, %broadcast.splat11
-
+  //
+  // Please note that the 'or' is equivalent to the 'and' here, this relies on
+  // BroadcastSplat being the IV which we know is a phi with 0 start and Lanes
+  // increment, which is all being checked below.
   Instruction *BroadcastSplat = nullptr;
   Constant *Const = nullptr;
   if (!match(TCP.Induction,
-             m_Add(m_Instruction(BroadcastSplat), m_Constant(Const))))
+             m_Add(m_Instruction(BroadcastSplat), m_Constant(Const))) &&
+      !match(TCP.Induction,
+             m_Or(m_Instruction(BroadcastSplat), m_Constant(Const))))
     return false;
 
   // Check that we're adding <0, 1, 2, 3...
   if (auto *CDS = dyn_cast<ConstantDataSequential>(Const)) {
     for (unsigned i = 0; i < CDS->getNumElements(); ++i) {
       if (CDS->getElementAsInteger(i) != i)
         return false;
     }

llvm/test/CodeGen/Thumb2/LowOverheadLoops/basic-tail-pred.ll

vector.ph:                                        ; preds = %entry
   call void @llvm.set.loop.iterations.i32(i32 %tmp13)
   br label %vector.body
 
 vector.body:                                      ; preds = %vector.body, %vector.ph
   %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
   %tmp14 = phi i32 [ %tmp13, %vector.ph ], [ %tmp15, %vector.body ]
   %broadcast.splatinsert = insertelement <16 x i32> undef, i32 %index, i32 0
   %broadcast.splat = shufflevector <16 x i32> %broadcast.splatinsert, <16 x i32> undef, <16 x i32> zeroinitializer
-  %induction = add <16 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %induction = or <16 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
   %tmp = getelementptr inbounds i8, i8* %a, i32 %index
   %tmp1 = icmp ule <16 x i32> %induction, %broadcast.splat11
   %tmp2 = bitcast i8* %tmp to <16 x i8>*
   %wide.masked.load = tail call <16 x i8> @llvm.masked.load.v16i8.p0v16i8(<16 x i8>* %tmp2, i32 4, <16 x i1> %tmp1, <16 x i8> undef)
   %tmp3 = getelementptr inbounds i8, i8* %b, i32 %index
   %tmp4 = bitcast i8* %tmp3 to <16 x i8>*
   %wide.masked.load2 = tail call <16 x i8> @llvm.masked.load.v16i8.p0v16i8(<16 x i8>* %tmp4, i32 4, <16 x i1> %tmp1, <16 x i8> undef)
   %mul = mul nsw <16 x i8> %wide.masked.load2, %wide.masked.load
vector.ph:                                        ; preds = %entry
   call void @llvm.set.loop.iterations.i32(i32 %tmp13)
   br label %vector.body
 
 vector.body:                                      ; preds = %vector.body, %vector.ph
   %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
   %tmp14 = phi i32 [ %tmp13, %vector.ph ], [ %tmp15, %vector.body ]
   %broadcast.splatinsert = insertelement <8 x i32> undef, i32 %index, i32 0
   %broadcast.splat = shufflevector <8 x i32> %broadcast.splatinsert, <8 x i32> undef, <8 x i32> zeroinitializer
   %induction = add <8 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   %tmp = getelementptr inbounds i16, i16* %a, i32 %index
   %tmp1 = icmp ule <8 x i32> %induction, %broadcast.splat11
   %tmp2 = bitcast i16* %tmp to <8 x i16>*
   %wide.masked.load = tail call <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>* %tmp2, i32 4, <8 x i1> %tmp1, <8 x i16> undef)
   %tmp3 = getelementptr inbounds i16, i16* %b, i32 %index
   %tmp4 = bitcast i16* %tmp3 to <8 x i16>*
   %wide.masked.load2 = tail call <8 x i16> @llvm.masked.load.v8i16.p0v8i16(<8 x i16>* %tmp4, i32 4, <8 x i1> %tmp1, <8 x i16> undef)
   %mul = mul nsw <8 x i16> %wide.masked.load2, %wide.masked.load
vector.ph:                                        ; preds = %entry
   call void @llvm.set.loop.iterations.i32(i32 %tmp13)
   br label %vector.body
 
 vector.body:                                      ; preds = %vector.body, %vector.ph
   %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
   %tmp14 = phi i32 [ %tmp13, %vector.ph ], [ %tmp15, %vector.body ]
   %broadcast.splatinsert = insertelement <4 x i32> undef, i32 %index, i32 0
   %broadcast.splat = shufflevector <4 x i32> %broadcast.splatinsert, <4 x i32> undef, <4 x i32> zeroinitializer
-  %induction = add <4 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3>
+  %induction = or <4 x i32> %broadcast.splat, <i32 0, i32 1, i32 2, i32 3>
   %tmp = getelementptr inbounds i32, i32* %a, i32 %index
   %tmp1 = icmp ule <4 x i32> %induction, %broadcast.splat11
   %tmp2 = bitcast i32* %tmp to <4 x i32>*
   %wide.masked.load = tail call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %tmp2, i32 4, <4 x i1> %tmp1, <4 x i32> undef)
   %tmp3 = getelementptr inbounds i32, i32* %b, i32 %index
   %tmp4 = bitcast i32* %tmp3 to <4 x i32>*
   %wide.masked.load2 = tail call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %tmp4, i32 4, <4 x i1> %tmp1, <4 x i32> undef)
   %mul = mul nsw <4 x i32> %wide.masked.load2, %wide.masked.load