[AArch64] Allow cost computation for interleaved accesses

[mgabka]

Summary

This commit is moving the rejection of scalable
  vectors in the loop vectorization up to the target specific cost model.
  Before this change for interleaved memory accesses there was an early
  rejection from inside the loop vectorizer cost model.

  The BasicTTIImplBase returns invalid cost for scalable vectors
  what should be enough for other backends to use.

  AArch64 will still return invalid cost for scalable vectors
  for interleaved accesses as more work is required to enable it.

[mgabka]

Looks like lack of vscale_range attribute was causing div by 0 exception.

[mgabka]

fix not closed attributes declaration in a test file

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14033–14034

The usage of getMinSVEVectorSizeInBits here is only relevant for fixed length vector types. When working with scalable types we already know the size of the legal types (i.e. vscale * 128) and so the : 128 part should be fine because Scalable_EC/(vscale * 128) ==> EC.getKnownMinValue() / 128).

Fixing this should mean you no longer require vscale_range attributes for your scalable vector tests.

14048–14049

Is this correct? Do you really want to say all scalable vector types are legal?

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

2489–2490

This is worth a comment because the naive optimisation is to simply move the check higher up whereas you deliberately want to exercise as much of the costing code as possible, even though we'll currently ignore the result.

2496–2498

Perhaps simplify to just "All other uses of scalable vectors are not legal."? With that said, looking at the base implementation of getInterleavedMemoryOpCost suggests it does the right thing so you could remove this code and just fall into the base version as before.

llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-accesses.ll

1–6 ↗

(On Diff #468552)

Does the code affect the output of these tests? Perhaps in this instance it is better to commit the tests first in isolation and then it's clear to see what effect this patch has (or rather, to show this patch does not change the existing behaviour).

[mgabka]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14033–14034

ok, I will simplify it

14048–14049

my understanding is that at the moment we do not know which scalable vector types are legal or not for this type of accesses, it depends how this is going to be implemented, right?

getInterleavedMemoryOpCost is just calling isLegalInterleavedAccessType, but later there is a check and returns invalid cost for all scalable VectorTypes.
I probably should add a TODO comment here.

I thought that this might be enough for now, we will need to add some extra tests later to check if the right types are legal or not.

At the moment the isLegalInterleavedAccessType is only used for the fixed width SVE so I think I should add extra assert here as well.

I might also try to check the scalable types here assuming legal element sizes etc, but at the moment we aren't able to test it.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

2489–2490

ok will add a comment, fair point

2496–2498

yeah you are right, it has an early check at the beginning, I will change it

llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-accesses.ll

1–6 ↗

(On Diff #468552)

My patch does not affect the output for this test.

This test is also a copy of llvm/test/Transforms/LoopVectorize/interleaved-accesses.ll
but uses aarch64 target triple and sve feature.
Is it ok to do it like that?

I wanted to demonstrate that despite removing the early exit from LoopVectorizationCostModel::getInterleaveGroupCost for scalable VF we still do not attempt to vectorize interleaved groups for AArch64.

In that instance yes It might be worth to pre commit this test.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14048–14049

This is target specific code so we can rely on implicit knowledge. For NEON you'll see the function only accepts the typical element types and then ends with Ensure the total vector size is 64 or a multiple of 128. So it seems reasonable to expect similar for scalable vectors.

llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-accesses.ll

1–6 ↗

(On Diff #468552)

Yep, by pre committing the test is becomes trivial to prove the patch works as expected by not changing the current behaviour.

[mgabka]

addressed review comments

[mgabka]

replaced by https://reviews.llvm.org/D145163

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

bool AArch64TargetLowering::hasPairedLoad(EVT LoadedType,
   return NumBits == 32 || NumBits == 64;
 }
 
 /// A helper function for determining the number of interleaved accesses we
 /// will generate when lowering accesses of the given type.
 unsigned AArch64TargetLowering::getNumInterleavedAccesses(
     VectorType *VecTy, const DataLayout &DL, bool UseScalable) const {
   unsigned VecSize = 128;
+  unsigned ElSize = DL.getTypeSizeInBits(VecTy->getElementType());
+  auto EC = VecTy->getElementCount();
   if (UseScalable)
     VecSize = std::max(Subtarget->getMinSVEVectorSizeInBits(), 128u);
-  return std::max<unsigned>(1, (DL.getTypeSizeInBits(VecTy) + 127) / VecSize);
+  return std::max<unsigned>(1,
+                            (EC.getKnownMinValue() * ElSize + 127) / VecSize);
 }
 
 MachineMemOperand::Flags
 AArch64TargetLowering::getTargetMMOFlags(const Instruction &I) const {
   if (Subtarget->getProcFamily() == AArch64Subtarget::Falkor &&
       I.getMetadata(FALKOR_STRIDED_ACCESS_MD) != nullptr)
     return MOStridedAccess;
   return MachineMemOperand::MONone;
 }
 
 bool AArch64TargetLowering::isLegalInterleavedAccessType(
     VectorType *VecTy, const DataLayout &DL, bool &UseScalable) const {
 
-  unsigned VecSize = DL.getTypeSizeInBits(VecTy);
   unsigned ElSize = DL.getTypeSizeInBits(VecTy->getElementType());
-  unsigned NumElements = cast<FixedVectorType>(VecTy)->getNumElements();
-
+  auto EC = VecTy->getElementCount();
   UseScalable = false;
 
   // Ensure that the predicate for this number of elements is available.
-  if (Subtarget->hasSVE() && !getSVEPredPatternFromNumElements(NumElements))
+  if (Subtarget->hasSVE() &&
+      !getSVEPredPatternFromNumElements(EC.getKnownMinValue()))
     return false;
 
   // Ensure the number of vector elements is greater than 1.
-  if (NumElements < 2)
+  if (EC.getKnownMinValue() < 2)
     return false;
 
   // Ensure the element type is legal.
   if (ElSize != 8 && ElSize != 16 && ElSize != 32 && ElSize != 64)
     return false;
 
+  if (EC.isScalable()) {
+    if (EC.getKnownMinValue() * ElSize == 128)
+      return true;
+    return false;
+  }
+
+  unsigned VecSize = DL.getTypeSizeInBits(VecTy);
   if (Subtarget->forceStreamingCompatibleSVE() ||
       (Subtarget->useSVEForFixedLengthVectors() &&
        (VecSize % Subtarget->getMinSVEVectorSizeInBits() == 0 ||
         (VecSize < Subtarget->getMinSVEVectorSizeInBits() &&
-         isPowerOf2_32(NumElements) && VecSize > 128)))) {
+         isPowerOf2_32(EC.getKnownMinValue()) && VecSize > 128)))) {
     UseScalable = true;
     return true;
   }
 
   // Ensure the total vector size is 64 or a multiple of 128. Types larger than
   // 128 will be split into multiple interleaved accesses.
   return VecSize == 64 || VecSize % 128 == 0;
 }

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

InstructionCost AArch64TTIImpl::getMemoryOpCost(unsigned Opcode, Type *Ty,
   return LT.first;
 }
 
 InstructionCost AArch64TTIImpl::getInterleavedMemoryOpCost(
     unsigned Opcode, Type *VecTy, unsigned Factor, ArrayRef<unsigned> Indices,
     Align Alignment, unsigned AddressSpace, TTI::TargetCostKind CostKind,
     bool UseMaskForCond, bool UseMaskForGaps) {
   assert(Factor >= 2 && "Invalid interleave factor");
-  auto *VecVTy = cast<FixedVectorType>(VecTy);
+  auto *VecVTy = cast<VectorType>(VecTy);
 
   if (!UseMaskForCond && !UseMaskForGaps &&
       Factor <= TLI->getMaxSupportedInterleaveFactor()) {
-    unsigned NumElts = VecVTy->getNumElements();
+    unsigned NumElts = VecVTy->getElementCount().getKnownMinValue();
     auto *SubVecTy =
-        FixedVectorType::get(VecTy->getScalarType(), NumElts / Factor);
+        VectorType::get(VecVTy->getElementType(),
+                        VecVTy->getElementCount().divideCoefficientBy(Factor));
 
     // ldN/stN only support legal vector types of size 64 or 128 in bits.
     // Accesses having vector types that are a multiple of 128 bits can be
     // matched to more than one ldN/stN instruction.
     bool UseScalable;
     if (NumElts % Factor == 0 &&
-        TLI->isLegalInterleavedAccessType(SubVecTy, DL, UseScalable))
-      return Factor * TLI->getNumInterleavedAccesses(SubVecTy, DL, UseScalable);
+        TLI->isLegalInterleavedAccessType(SubVecTy, DL, UseScalable)) {
+      unsigned Cost =
+          Factor * TLI->getNumInterleavedAccesses(SubVecTy, DL, UseScalable);
+      // Deliberately do not move it to the begining of this function
+      // as we want to execute as much as possible code for scalable vectors.
+      if (isa<ScalableVectorType>(VecTy))
+        return InstructionCost::getInvalid();
+      return Cost;
+    }
   }
 
   return BaseT::getInterleavedMemoryOpCost(Opcode, VecTy, Factor, Indices,
                                            Alignment, AddressSpace, CostKind,
                                            UseMaskForCond, UseMaskForGaps);
 }
 
 InstructionCost
 AArch64TTIImpl::getCostOfKeepingLiveOverCall(ArrayRef<Type *> Tys) {
   InstructionCost Cost = 0;
   TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
   for (auto *I : Tys) {

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

return TTI.getAddressComputationCost(VectorTy) +
          TTI.getGatherScatterOpCost(
              I->getOpcode(), VectorTy, Ptr, Legal->isMaskRequired(I), Alignment,
              TargetTransformInfo::TCK_RecipThroughput, I);
 }
 
 InstructionCost
 LoopVectorizationCostModel::getInterleaveGroupCost(Instruction *I,
                                                    ElementCount VF) {
-  // TODO: Once we have support for interleaving with scalable vectors
-  // we can calculate the cost properly here.
-  if (VF.isScalable())
-    return InstructionCost::getInvalid();
-
   Type *ValTy = getLoadStoreType(I);
   auto *VectorTy = cast<VectorType>(ToVectorTy(ValTy, VF));
   unsigned AS = getLoadStoreAddressSpace(I);
   enum TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
 
   auto Group = getInterleavedAccessGroup(I);
   assert(Group && "Fail to get an interleaved access group.");