[SVE][LoopVectorize] Add masked load/store and gather/scatter support for SVE

[david-arm]

Summary

This patch updates IRBuilder::CreateMaskedGather/Scatter to work
with ScalableVectorType and adds isLegalMaskedGather/Scatter functions
to AArch64TargetTransformInfo. In addition I've fixed up
isLegalMaskedLoad/Store to return true for supported scalar types,
since this is what the vectorizer asks for.

In LoopVectorize.cpp I've changed
LoopVectorizationCostModel::getInterleaveGroupCost to return an invalid
cost for scalable vectors, since currently this relies upon using shuffle
vector for reversing vectors. In addition, in
LoopVectorizationCostModel::setCostBasedWideningDecision I have assumed
that the cost of scalarising memory ops is infinitely expensive.

I have added some simple masked load/store and gather/scatter tests,
including cases where we use gathers and scatters for conditional invariant
loads and stores.

[david-arm]

• Only bail out of calculating group interleave costs if this is a reversed group.
• Fixed some formatting issues.

[david-arm]

• Changed setCostBasedWideningDecision to use InstructionCost::getInvalid() when the cost is invalid.
• Changed getInterleaveGroupCost to return an invalid cost for SVE.

[fhahn]

llvm/test/Transforms/LoopVectorize/AArch64/sve-gather-scatter.ll

13

This shouldn't be necessary, LV will generate a minimum iteration check anyways.

16

are those blocks needed?

29

Is this necessary? Can we instead just load from a i64* or use i32 in the GEP? Same for the other tests.

[david-arm]

llvm/test/Transforms/LoopVectorize/AArch64/sve-gather-scatter.ll

13

Sure, I simply adapted these tests by dumping out the module before the vectoriser, i.e. I didn't hand craft the tests. :) I can remove them no problem.

16

Probably not. Again I didn't write these IR tests by hand - I dumped out the module before the vectoriser so that explains why these blocks are here. I can try to clean them up.

29

I deliberately left this in as the sext exposes a different code path in the vectoriser that was broken - see line 7463 in LoopVectorize.cpp. I'd like to have at least one test with a sext of i32 -> i64 so that the fix and test are in the same patch.

[david-arm]

• Tried to tidy up the tests a bit and remove unnecessary blocks.

[fhahn]

llvm/test/Transforms/LoopVectorize/AArch64/sve-gather-scatter.ll

13

Thanks. This will make the life of anyone who has to adjust the tests in the future easier.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7081

I just realised this TODO can probably be removed, which I can do before merging if that's ok?

[kmclaughlin]

LGTM!

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

221

nit: I think this is the same check as we have at the top of isLegalMaskedLoadStore - could one of them maybe be changed so that they are consistent?

[david-arm]

• After @kmclaughlin's suggestion to unify the behaviour between the masked loads/stores and gathers/scatters I realised that we should be checking for scalar types isLegalMaskedLoad/Store in the same way as isLegalMaskedGather/Scatter does. This is what other targets do because the vectorizer queries support based on the scalar type.
• Added tests that we vectorise with masked loads and stores.

[david-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

221

That's a cracking bug you found @kmclaughlin! It turns out we weren't properly supporting masked loads and stores in the vectoriser either. Thanks!

[kmclaughlin]

Thanks for making these changes @david-arm, LGTM

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 NumElts = PtrsTy->getElementCount();
+  auto *DataTy = VectorType::get(PtrTy->getElementType(), NumElts);
 
   if (!Mask)
     Mask = Constant::getAllOnesValue(
-        FixedVectorType::get(Type::getInt1Ty(Context), NumElts));
+        VectorType::get(Type::getInt1Ty(Context), NumElts));
 
   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 NumElts = PtrsTy->getElementCount();
 
 #ifndef NDEBUG
   auto PtrTy = cast<PointerType>(PtrsTy->getElementType());
-  assert(NumElts == DataTy->getNumElements() &&
+  assert(NumElts == 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), NumElts));
 
   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/Target/AArch64/AArch64TargetTransformInfo.h

public:
   void getPeelingPreferences(Loop *L, ScalarEvolution &SE,
                              TTI::PeelingPreferences &PP);
 
   Value *getOrCreateResultFromMemIntrinsic(IntrinsicInst *Inst,
                                            Type *ExpectedType);
 
   bool getTgtMemIntrinsic(IntrinsicInst *Inst, MemIntrinsicInfo &Info);
 
-  bool isLegalMaskedLoadStore(Type *DataType, Align Alignment) {
-    if (!isa<ScalableVectorType>(DataType) || !ST->hasSVE())
-      return false;
-
-    Type *Ty = cast<ScalableVectorType>(DataType)->getElementType();
+  bool isLegalScalarTypeForSVEMaskedMemOp(Type *Ty) const {
     if (Ty->isPointerTy())
       return true;
 
     if (Ty->isBFloatTy() || Ty->isHalfTy() ||
         Ty->isFloatTy() || Ty->isDoubleTy())
       return true;
 
     if (Ty->isIntegerTy(8) || Ty->isIntegerTy(16) ||
         Ty->isIntegerTy(32) || Ty->isIntegerTy(64))
       return true;
 
     return false;
   }
 
+  bool isLegalMaskedLoadStore(Type *DataType, Align Alignment) {
+    if (!isa<ScalableVectorType>(DataType) || !ST->hasSVE())
+      return false;
+
+    Type *Ty = cast<ScalableVectorType>(DataType)->getElementType();
+    return isLegalScalarTypeForSVEMaskedMemOp(Ty);
+  }
+
   bool isLegalMaskedLoad(Type *DataType, Align Alignment) {
     return isLegalMaskedLoadStore(DataType, Alignment);
   }
 
   bool isLegalMaskedStore(Type *DataType, Align Alignment) {
     return isLegalMaskedLoadStore(DataType, Alignment);
   }
 
+  bool isLegalMaskedGatherScatter(Type *DataType) const {
+    if (isa<FixedVectorType>(DataType) || !ST->hasSVE())
+      return false;
+
+    return isLegalScalarTypeForSVEMaskedMemOp(DataType->getScalarType());
+  }
+
+  bool isLegalMaskedGather(Type *DataType, Align Alignment) const {
+    return isLegalMaskedGatherScatter(DataType);
+  }
+  bool isLegalMaskedScatter(Type *DataType, Align Alignment) const {
+    return isLegalMaskedGatherScatter(DataType);
+  }
+
   bool isLegalNTStore(Type *DataType, Align Alignment) {
     // NOTE: The logic below is mostly geared towards LV, which calls it with
     //       vectors with 2 elements. We might want to improve that, if other
     //       users show up.
     // Nontemporal vector stores can be directly lowered to STNP, if the vector
     // can be halved so that each half fits into a register. That's the case if
     // the element type fits into a register and the number of elements is a
     // power of 2 > 1.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

for (Instruction &I : *BB) {
         auto Group = getInterleavedAccessGroup(&I);
         assert(Group && "Fail to get an interleaved access group.");
 
         // Make one decision for the whole group.
         if (getWideningDecision(&I, VF) != CM_Unknown)
           continue;
 
         NumAccesses = Group->getNumMembers();
-        if (interleavedAccessCanBeWidened(&I, VF))
+        // TODO: For now we don't consider scalable vectors here as it
+        // potentially requires shuffle support for reversing vectors.
+        if (interleavedAccessCanBeWidened(&I, VF) &&
+            !VF.isScalable())
           InterleaveCost = getInterleaveGroupCost(&I, VF);
       }
 
       InstructionCost GatherScatterCost =
           isLegalGatherOrScatter(&I)
               ? getGatherScatterCost(&I, VF) * NumAccesses
               : std::numeric_limits<int>::max();
 
-      InstructionCost ScalarizationCost =
-          getMemInstScalarizationCost(&I, VF) * NumAccesses;
+      InstructionCost ScalarizationCost = !VF.isScalable() ?
+          getMemInstScalarizationCost(&I, VF) * NumAccesses
+          : std::numeric_limits<int>::max();
 
       // Choose better solution for the current VF,
       // write down this decision and use it during vectorization.
       InstructionCost Cost;
       InstWidening Decision;
       if (InterleaveCost <= GatherScatterCost &&
           InterleaveCost < ScalarizationCost) {
         Decision = CM_Interleave;
         Cost = InterleaveCost;
       } else if (GatherScatterCost < ScalarizationCost) {
         Decision = CM_GatherScatter;
         Cost = GatherScatterCost;
       } else {
+        assert(!VF.isScalable() && "We cannot yet scalarise for scalable vectors");
         Decision = CM_Scalarize;
         Cost = ScalarizationCost;
       }
       // If the instructions belongs to an interleave group, the whole group
       // receives the same decision. The whole group receives the cost, but
       // the cost will actually be assigned to one instruction.
       if (auto Group = getInterleavedAccessGroup(&I))
         setWideningDecision(Group, VF, Decision, Cost);
if (canTruncateToMinimalBitwidth(I, VF)) {
             largestIntegerVectorType(ToVectorTy(I->getType(), VF), MinVecTy);
       } else if (Opcode == Instruction::ZExt || Opcode == Instruction::SExt) {
         SrcVecTy = largestIntegerVectorType(SrcVecTy, MinVecTy);
         VectorTy =
             smallestIntegerVectorType(ToVectorTy(I->getType(), VF), MinVecTy);
       }
     }
 
+    unsigned N;
+    if (isScalarAfterVectorization(I, VF)) {
       assert(!VF.isScalable() && "VF is assumed to be non scalable");
-    unsigned N = isScalarAfterVectorization(I, VF) ? VF.getKnownMinValue() : 1;
+      N = VF.getKnownMinValue();
+    } else
+      N = 1;
     return N *
            TTI.getCastInstrCost(Opcode, VectorTy, SrcVecTy, CCH, CostKind, I);
   }
   case Instruction::Call: {
     bool NeedToScalarize;
     CallInst *CI = cast<CallInst>(I);
     InstructionCost CallCost = getVectorCallCost(CI, VF, NeedToScalarize);
     if (getVectorIntrinsicIDForCall(CI, TLI)) {

llvm/test/Transforms/LoopVectorize/AArch64/sve-gather-scatter.ll

+; RUN: opt -loop-vectorize -dce -instcombine -mtriple aarch64-linux-gnu -mattr=+sve -S %s -o - | FileCheck %s
+
+define void @gather_nxv4i32(float* noalias nocapture readonly %a, i32* noalias nocapture readonly %b, float* noalias nocapture %c, i64 %n) {
+; CHECK-LABEL: @gather_nxv4i32
+; CHECK: vector.body:
+; CHECK:   %[[IND:.*]] = load <vscale x 4 x i32>, <vscale x 4 x i32>*
+; CHECK:   %[[EXTIND:.*]] = sext <vscale x 4 x i32> %[[IND]] to <vscale x 4 x i64>
+; CHECK:   %[[PTRS:.*]] = getelementptr inbounds float, float* %a, <vscale x 4 x i64> %[[EXTIND]]
+; CHECK:   %[[GLOAD:.*]] = call <vscale x 4 x float> @llvm.masked.gather.nxv4f32.nxv4p0f32(<vscale x 4 x float*> %[[PTRS]]
+; CHECK:   store <vscale x 4 x float> %[[GLOAD]], <vscale x 4 x float>*
+entry:
+  %cmp10 = icmp sgt i64 %n, 0
+  br i1 %cmp10, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %entry
+  br label %for.body
+
+for.cond.cleanup.loopexit:                        ; preds = %for.body
+  br label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %for.body.preheader ]
+  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %idxprom2 = sext i32 %0 to i64
+  %arrayidx3 = getelementptr inbounds float, float* %a, i64 %idxprom2
+  %1 = load float, float* %arrayidx3, align 4
+  %arrayidx5 = getelementptr inbounds float, float* %c, i64 %indvars.iv
+  store float %1, float* %arrayidx5, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond.not, label %for.cond.cleanup.loopexit, label %for.body, !llvm.loop !0
+}
+
+define void @scatter_nxv4i32(float* noalias nocapture %a, i32* noalias nocapture readonly %b, float* noalias nocapture readonly %c, i64 %n) {
+; CHECK-LABEL: @scatter_nxv4i32
+; CHECK: vector.body:
+; CHECK:   %[[VALS:.*]] = load <vscale x 4 x float>
+; CHECK:   %[[IND:.*]] = load <vscale x 4 x i32>, <vscale x 4 x i32>* %7, align 4
+; CHECK:   %[[EXTIND:.*]] = sext <vscale x 4 x i32> %[[IND]] to <vscale x 4 x i64>
+; CHECK:   %[[PTRS:.*]] = getelementptr inbounds float, float* %a, <vscale x 4 x i64> %[[EXTIND]]
+; CHECK:   call void @llvm.masked.scatter.nxv4f32.nxv4p0f32(<vscale x 4 x float> %[[VALS]], <vscale x 4 x float*> %[[PTRS]]
+entry:
+  %cmp10 = icmp sgt i64 %n, 0
+  br i1 %cmp10, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+
+for.body:                                         ; preds = %entry, %for.body
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+  %arrayidx = getelementptr inbounds float, float* %c, i64 %indvars.iv
+  %0 = load float, float* %arrayidx, align 4
+  %arrayidx3 = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %1 = load i32, i32* %arrayidx3, align 4
+  %idxprom4 = sext i32 %1 to i64
+  %arrayidx5 = getelementptr inbounds float, float* %a, i64 %idxprom4
+  store float %0, float* %arrayidx5, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !0
+}
+
+define void @scatter_inv_nxv4i32(i32* noalias nocapture %inv, i32* noalias nocapture readonly %b, i64 %n) {
+; CHECK-LABEL: @scatter_inv_nxv4i32
+; CHECK: vector.ph:
+; CHECK:   %[[INS:.*]] = insertelement <vscale x 4 x i32*> poison, i32* %inv, i32 0
+; CHECK:   %[[PTRSPLAT:.*]] = shufflevector <vscale x 4 x i32*> %[[INS]], <vscale x 4 x i32*> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK: vector.body:
+; CHECK:   %[[VALS:.*]] = load <vscale x 4 x i32>, <vscale x 4 x i32>* %5, align 4
+; CHECK:   %[[MASK:.*]] = icmp ne <vscale x 4 x i32> %[[VALS]],
+; CHECK:   call void @llvm.masked.scatter.nxv4i32.nxv4p0i32({{.*}}, <vscale x 4 x i32*> %[[PTRSPLAT]], i32 4, <vscale x 4 x i1> %[[MASK]])
+entry:
+  %cmp5 = icmp sgt i64 %n, 0
+  br i1 %cmp5, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %entry
+  br label %for.body
+
+for.cond.cleanup.loopexit:                        ; preds = %for.inc
+  br label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.inc
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.inc ], [ 0, %for.body.preheader ]
+  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %tobool.not = icmp eq i32 %0, 0
+  br i1 %tobool.not, label %for.inc, label %if.then
+
+if.then:                                          ; preds = %for.body
+  store i32 3, i32* %inv, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body, %if.then
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond.not, label %for.cond.cleanup.loopexit, label %for.body, !llvm.loop !0
+}
+
+define void @gather_inv_nxv4i32(i32* noalias nocapture %a, i32* noalias nocapture readonly %inv, i64 %n) {
+; CHECK-LABEL: @gather_inv_nxv4i32
+; CHECK: vector.ph:
+; CHECK:   %[[INS:.*]] = insertelement <vscale x 4 x i32*> poison, i32* %inv, i32 0
+; CHECK:   %[[PTRSPLAT:.*]] = shufflevector <vscale x 4 x i32*> %[[INS]], <vscale x 4 x i32*> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK: vector.body:
+; CHECK:   %[[VALS:.*]] = load <vscale x 4 x i32>, <vscale x 4 x i32>* %5, align 4
+; CHECK:   %[[MASK:.*]] = icmp sgt <vscale x 4 x i32> %[[VALS]],
+; CHECK:   %{{.*}} = call <vscale x 4 x i32> @llvm.masked.gather.nxv4i32.nxv4p0i32(<vscale x 4 x i32*> %[[PTRSPLAT]], i32 4, <vscale x 4 x i1> %[[MASK]]
+entry:
+  %cmp11 = icmp sgt i64 %n, 0
+  br i1 %cmp11, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.inc, %entry
+  ret void
+
+for.body:                                         ; preds = %entry, %for.inc
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.inc ], [ 0, %entry ]
+  %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %cmp2 = icmp sgt i32 %0, 3
+  br i1 %cmp2, label %if.then, label %for.inc
+
+if.then:                                          ; preds = %for.body
+  %1 = load i32, i32* %inv, align 4
+  store i32 %1, i32* %arrayidx, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body, %if.then
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !0
+}
+
+!0 = distinct !{!0, !1, !2, !3, !4, !5}
+!1 = !{!"llvm.loop.mustprogress"}
+!2 = !{!"llvm.loop.vectorize.width", i32 4}
+!3 = !{!"llvm.loop.vectorize.scalable.enable", i1 true}
+!4 = !{!"llvm.loop.interleave.count", i32 1}
+!5 = !{!"llvm.loop.vectorize.enable", i1 true}