[SVE] In LoopIdiomRecognize::isLegalStore bail out for scalable vectors

[david-arm]

Summary

The function LoopIdiomRecognize::isLegalStore looks for stores in loops
that could be transformed into memset or memcpy. However, the algorithm
currently requires that we know how big the store is at runtime, i.e.
that the store size will not overflow an unsigned integer. For scalable
vectors we cannot guarantee this so I have changed the code to bail out
for now. In addition, even if we add a way to query the maximum value of
vscale in future we will still need to update the algorithm to cope with
non-constant strides. The additional cost associated with calculating
the memset and memcpy arguments will need to be taken into account as
well.

This patch also fixes up an implicit TypeSize -> uint64_t cast,
thereby removing a warning. I've added tests here showing a fixed
width vector loop being transformed into memcpy, and a scalable
vector loop remaining unchanged:

Transforms/LoopIdiom/memcpy-vectors.ll

[efriedma]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

475

Do you have any idea why we're checking whether the store overflows an unsigned?

[david-arm]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

475

Hi @efriedma, honestly I don't know why we check for overflow here, although I assumed it was probably because they wanted to multiply the store size by something else at some point, so want to ensure it fits in a 64-bit value. That's a pure guess though! I can look into this further if you want, but there is still the issue of non-constant strides and extra maths to deal with.

[efriedma]

LGTM

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

475

In that case, please simplify the comment to just say the code needs changes to deal with a non-constant stride.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

LoopIdiomRecognize::isLegalStore(StoreInst *SI) {
   // Avoid merging nontemporal stores.
   if (SI->getMetadata(LLVMContext::MD_nontemporal))
     return LegalStoreKind::None;
 
   Value *StoredVal = SI->getValueOperand();
   Value *StorePtr = SI->getPointerOperand();
 
   // Reject stores that are so large that they overflow an unsigned.
-  uint64_t SizeInBits = DL->getTypeSizeInBits(StoredVal->getType());
-  if ((SizeInBits & 7) || (SizeInBits >> 32) != 0)
+  // When storing out scalable vectors we bail out for now, since the code
+  // below currently only works for constant strides.
+  TypeSize SizeInBits = DL->getTypeSizeInBits(StoredVal->getType());
+  if (SizeInBits.isScalable() || (SizeInBits.getFixedSize() & 7) ||
+      (SizeInBits.getFixedSize() >> 32) != 0)
     return LegalStoreKind::None;
 
   // See if the pointer expression is an AddRec like {base,+,1} on the current
   // loop, which indicates a strided store.  If we have something else, it's a
   // random store we can't handle.
   const SCEVAddRecExpr *StoreEv =
       dyn_cast<SCEVAddRecExpr>(SE->getSCEV(StorePtr));
   if (!StoreEv || StoreEv->getLoop() != CurLoop || !StoreEv->isAffine())

llvm/test/Transforms/LoopIdiom/memcpy-vectors.ll

+; RUN: opt -loop-idiom -S <%s 2>%t | FileCheck %s
+; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
+
+; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
+; WARN-NOT: warning
+
+define void @memcpy_fixed_vec(i64* noalias %a, i64* noalias %b) local_unnamed_addr #1 {
+; CHECK-LABEL: @memcpy_fixed_vec(
+; CHECK: entry:
+; CHECK: memcpy
+; CHECK: vector.body
+entry:
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %entry
+  %index = phi i64 [ 0, %entry ], [ %index.next, %vector.body ]
+  %0 = getelementptr inbounds i64, i64* %a, i64 %index
+  %1 = bitcast i64* %0 to <2 x i64>*
+  %wide.load = load <2 x i64>, <2 x i64>* %1, align 8
+  %2 = getelementptr inbounds i64, i64* %b, i64 %index
+  %3 = bitcast i64* %2 to <2 x i64>*
+  store <2 x i64> %wide.load, <2 x i64>* %3, align 8
+  %index.next = add nuw nsw i64 %index, 2
+  %4 = icmp eq i64 %index.next, 1024
+  br i1 %4, label %for.cond.cleanup, label %vector.body
+
+for.cond.cleanup:                                 ; preds = %vector.body
+  ret void
+}
+
+define void @memcpy_scalable_vec(i64* noalias %a, i64* noalias %b) local_unnamed_addr #1 {
+; CHECK-LABEL: @memcpy_scalable_vec(
+; CHECK: entry:
+; CHECK-NOT: memcpy
+; CHECK: vector.body
+entry:
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %entry
+  %index = phi i64 [ 0, %entry ], [ %index.next, %vector.body ]
+  %0 = bitcast i64* %a to <vscale x 2 x i64>*
+  %1 = getelementptr inbounds <vscale x 2 x i64>, <vscale x 2 x i64>* %0, i64 %index
+  %wide.load = load <vscale x 2 x i64>, <vscale x 2 x i64>* %1, align 16
+  %2 = bitcast i64* %b to <vscale x 2 x i64>*
+  %3 = getelementptr inbounds <vscale x 2 x i64>, <vscale x 2 x i64>* %2, i64 %index
+  store <vscale x 2 x i64> %wide.load, <vscale x 2 x i64>* %3, align 16
+  %index.next = add nuw nsw i64 %index, 1
+  %4 = icmp eq i64 %index.next, 1024
+  br i1 %4, label %for.cond.cleanup, label %vector.body
+
+for.cond.cleanup:                                 ; preds = %vector.body
+  ret void
+}