[POC][LoopVectorizer] Allow invariant loads/stores using masked gather/scatter for a scalable VF.

[sdesmalen]

Summary

This patch is part of a proof of concept for vectorising a loop using
scalable vectors. The patch is shared for reference and there is no
expectation for this patch to land in the current form.

For fixed-width vectors, the loopvectorizer assumes that certain operations
can be scalarized. For example, loads/stores from uniform pointers without
masking are scalarized, which is not possible for scalable vectors. For
these, use gather/scatter instructions instead until we've found a way to
properly widen these types.

void loop(int N, double *a, double *b) {
  #pragma clang loop vectorize_width(4, scalable)
  for (int i = 0; i < N; i++) {
    a[42] = b[i] + 1.0;   // uses llvm.masked.scatter for the store
  }   
}

[ctetreau]

llvm/test/Transforms/LoopVectorize/scalable-loop-invariant-store-unpredicated-body-scalar-tail.ll

3

I suppose this is why you don't want to actually merge this currently? What happens if if it's not aarch64?

[sdesmalen]

llvm/test/Transforms/LoopVectorize/scalable-loop-invariant-store-unpredicated-body-scalar-tail.ll

3

The REQUIRES: aarch64-registered-target is actually unnecessary, not sure why I thought this was needed.
This patch is probably simple enough to be reviewed as-is. The other POC patches I've split up into smaller NFC patches, but there is little to simplify here.

[ctetreau]

I have no objections to this, but you should probably get some more eyes on it.

[sdesmalen]

This has since been superseded by other patches.

llvm/lib/IR/IRBuilder.cpp

 /// \p Mask     - vector of booleans which indicates what vector lanes should
 ///               be accessed in memory
 /// \p PassThru - pass-through value that is used to fill the masked-off lanes
 ///               of the result
 /// \p Name     - name of the result variable
 CallInst *IRBuilderBase::CreateMaskedGather(Value *Ptrs, Align Alignment,
                                             Value *Mask, Value *PassThru,
                                             const Twine &Name) {
-  auto *PtrsTy = cast<FixedVectorType>(Ptrs->getType());
+  auto *PtrsTy = cast<VectorType>(Ptrs->getType());
   auto *PtrTy = cast<PointerType>(PtrsTy->getElementType());
-  unsigned NumElts = PtrsTy->getNumElements();
-  auto *DataTy = FixedVectorType::get(PtrTy->getElementType(), NumElts);
+  ElementCount EC = PtrsTy->getElementCount();
+  auto *DataTy = VectorType::get(PtrTy->getElementType(), EC);
 
   if (!Mask)
     Mask = Constant::getAllOnesValue(
-        FixedVectorType::get(Type::getInt1Ty(Context), NumElts));
+        VectorType::get(Type::getInt1Ty(Context), EC));
 
   if (!PassThru)
     PassThru = UndefValue::get(DataTy);
 
   Type *OverloadedTypes[] = {DataTy, PtrsTy};
   Value *Ops[] = {Ptrs, getInt32(Alignment.value()), Mask, PassThru};
 
   // We specify only one type when we create this intrinsic. Types of other
   // arguments are derived from this type.
   return CreateMaskedIntrinsic(Intrinsic::masked_gather, Ops, OverloadedTypes,
                                Name);
 }
 
 /// Create a call to a Masked Scatter intrinsic.
 /// \p Data  - data to be stored,
 /// \p Ptrs  - the vector of pointers, where the \p Data elements should be
 ///            stored
 /// \p Align - alignment for one element
 /// \p Mask  - vector of booleans which indicates what vector lanes should
 ///            be accessed in memory
 CallInst *IRBuilderBase::CreateMaskedScatter(Value *Data, Value *Ptrs,
                                              Align Alignment, Value *Mask) {
-  auto *PtrsTy = cast<FixedVectorType>(Ptrs->getType());
-  auto *DataTy = cast<FixedVectorType>(Data->getType());
-  unsigned NumElts = PtrsTy->getNumElements();
+  auto *PtrsTy = cast<VectorType>(Ptrs->getType());
+  auto *DataTy = cast<VectorType>(Data->getType());
+  ElementCount EC = PtrsTy->getElementCount();
 
 #ifndef NDEBUG
   auto PtrTy = cast<PointerType>(PtrsTy->getElementType());
-  assert(NumElts == DataTy->getNumElements() &&
+  assert(EC == DataTy->getElementCount() &&
          PtrTy->getElementType() == DataTy->getElementType() &&
          "Incompatible pointer and data types");
 #endif
 
   if (!Mask)
     Mask = Constant::getAllOnesValue(
-        FixedVectorType::get(Type::getInt1Ty(Context), NumElts));
+        VectorType::get(Type::getInt1Ty(Context), EC));
 
   Type *OverloadedTypes[] = {DataTy, PtrsTy};
   Value *Ops[] = {Data, Ptrs, getInt32(Alignment.value()), Mask};
 
   // We specify only one type when we create this intrinsic. Types of other
   // arguments are derived from this type.
   return CreateMaskedIntrinsic(Intrinsic::masked_scatter, Ops, OverloadedTypes);
 }

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

for (Instruction &I : *BB) {
           // Conditional loads and stores should be scalarized and predicated.
           // isScalarWithPredication cannot be used here since masked
           // gather/scatters are not considered scalar with predication.
           !Legal->blockNeedsPredication(I.getParent())) {
         // TODO: Avoid replicating loads and stores instead of
         // relying on instcombine to remove them.
         // Load: Scalar load + broadcast
         // Store: Scalar store + isLoopInvariantStoreValue ? 0 : extract
+
+        // For Scalable vectors, use G/S instructions instead of a scalarised
+        // store for now, until we add a mechanism to extract the last lane of
+        // the vector. This is allowed, since llvm.masked.scatter is guaranteed
+        // to be an ordered store (from first to last element).
+        if (VF.isScalable())
+          setWideningDecision(&I, VF, CM_GatherScatter,
+                              /*Cost=*/VF.getKnownMinValue());
+        else {
           unsigned Cost = getUniformMemOpCost(&I, VF);
           setWideningDecision(&I, VF, CM_Scalarize, Cost);
+        }
         continue;
       }
 
       // We assume that widening is the best solution when possible.
       if (memoryInstructionCanBeWidened(&I, VF)) {
         unsigned Cost = getConsecutiveMemOpCost(&I, VF);
         int ConsecutiveStride =
                Legal->isConsecutivePtr(getLoadStorePointerOperand(&I));
         assert((ConsecutiveStride == 1 || ConsecutiveStride == -1) &&
                "Expected consecutive stride.");
         InstWidening Decision =
             ConsecutiveStride == 1 ? CM_Widen : CM_Widen_Reverse;
         setWideningDecision(&I, VF, Decision, Cost);
         continue;
       }
 
+      // For Scalable vectors, all nodes must be widened using gather/scatter
+      // or interleaved loads/stores. For now, default to gather/scatter.
+      if (VF.isScalable()) {
+        setWideningDecision(&I, VF, CM_GatherScatter,
+                            /*Cost=*/VF.getKnownMinValue());
+        continue;
+      }
+
       // Choose between Interleaving, Gather/Scatter or Scalarization.
       unsigned InterleaveCost = std::numeric_limits<unsigned>::max();
       unsigned NumAccesses = 1;
       if (isAccessInterleaved(&I)) {
         auto Group = getInterleavedAccessGroup(&I);
         assert(Group && "Fail to get an interleaved access group.");
 
         // Make one decision for the whole group.

llvm/test/Transforms/LoopVectorize/scalable-loop-invariant-store-unpredicated-body-scalar-tail.ll

+; For now this test requires aarch64-registered-target, until we can
+; also pass the loop hint as a 'force-vector-width' flag to opt.
+; REQUIRES: aarch64-registered-target
+
+; RUN: opt -S -loop-vectorize -instcombine < %s | FileCheck %s
+
+source_filename = "loop.c"
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+; This test has an invariant store to a[42], which it vectorizes using a masked.scatter operation.
+
+; CHECK: for.body.preheader:
+; CHECK-DAG: %wide.trip.count = zext i32 %N to i64
+; CHECK-DAG: %[[VSCALE:.*]] = call i64 @llvm.vscale.i64()
+; CHECK-DAG: %[[VSCALEX4:.*]] = shl i64 %[[VSCALE]], 2
+; CHECK-DAG: %min.iters.check = icmp ugt i64 %[[VSCALEX4]], %wide.trip.count
+
+; CHECK: vector.ph:
+; CHECK-DAG:  %[[VSCALE:.*]] = call i64 @llvm.vscale.i64()
+; CHECK-DAG:  %[[VSCALEX4:.*]] = shl i64 %[[VSCALE]], 2
+; CHECK-DAG:  %n.mod.vf = urem i64 %wide.trip.count, %[[VSCALEX4]]
+; CHECK:      %n.vec = sub nsw i64 %wide.trip.count, %n.mod.vf
+
+; CHECK: vector.body:
+; CHECK: %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+; CHECK: %[[IDXB:.*]] = getelementptr inbounds double, double* %b, i64 %index
+; CHECK: %[[IDXB_CAST:.*]] = bitcast double* %[[IDXB]] to <vscale x 4 x double>*
+; CHECK: %wide.load = load <vscale x 4 x double>, <vscale x 4 x double>* %[[IDXB_CAST]], align 8, !alias.scope !0
+; CHECK: %[[FADD:.*]] = fadd <vscale x 4 x double> %wide.load, shufflevector (<vscale x 4 x double> insertelement (<vscale x 4 x double> undef, double 1.000000e+00, i32 0), <vscale x 4 x double> undef, <vscale x 4 x i32> zeroinitializer)
+; CHECK: %[[IDXA:.*]] = getelementptr inbounds double, double* %a, i64 42
+; CHECK: %[[SPLATINSERT:.*]] = insertelement <vscale x 4 x double*> undef, double* %[[IDXA]], i32 0
+; CHECK: %[[SPLAT:.*]] = shufflevector <vscale x 4 x double*> %[[SPLATINSERT]], <vscale x 4 x double*> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK: call void @llvm.masked.scatter.nxv4f64.nxv4p0f64(<vscale x 4 x double> %[[FADD]], <vscale x 4 x double*> %[[SPLAT]], i32 8, <vscale x 4 x i1> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> undef, i1 true, i32 0), <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer))
+; CHECK: %[[VSCALE:.*]] = call i64 @llvm.vscale.i64()
+; CHECK: %[[VSCALEX4:.*]] = shl i64 %[[VSCALE]], 2
+; CHECK: %index.next = add i64 %index, %[[VSCALEX4]]
+; CHECK: %[[CMP:.*]] = icmp eq i64 %index.next, %n.vec
+; CHECK: br i1 %[[CMP]], label %middle.block, label %vector.body, !llvm.loop !5
+
+define void @loop(i32 %N, double* nocapture %a, double* nocapture readonly %b) {
+entry:
+  %cmp7 = icmp sgt i32 %N, 0
+  br i1 %cmp7, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %N to i64
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds double, double* %b, i64 %indvars.iv
+  %0 = load double, double* %arrayidx, align 8
+  %add = fadd double %0, 1.000000e+00
+  %arrayidx2 = getelementptr inbounds double, double* %a, i64 42
+  store double %add, double* %arrayidx2, align 8
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !1
+}
+
+!1 = distinct !{!1, !2, !3}
+!2 = !{!"llvm.loop.vectorize.width", !4}
+!3 = !{!"llvm.loop.interleave.count", i32 1}
+!4 = !{i32 4, i1 1}