[LoopVectorize] Don't attempt to widen certain calls for scalable vectors

[david-arm]

Summary

When attempting to vectorise a loop containing a call to the following
intrinsic

tail call fast float @llvm.sin.f32(float %0)

we then crashed in tryToWidenCall because we don't have scalable vector
versions of these intrinsics and we also cannot scalarise them. Rather
than rely on the cost model returning Invalid for such cases we should
be explicitly rejecting scalable vectorisation upfront.

This patch adds a new TTI interface called isLegalVectorIntrinsic that
checks if it is legal to vectorise a particular intrinsic. We then
walk through all calls in the loop ensuring they are legal - if not
we return a zero VF from LoopVectorizationCostModel::getMaxLegalScalableVF.

Tests have been added here:

Transforms/LoopVectorize/AArch64/scalable-call.ll

[kmclaughlin]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5654

nit: could this be moved below the comments beneath it so that it's above if !(canVectorizeReductions(VF))?

5672

Hi @david-arm, is it possible for there to be a VecID but getMappings(*CI) is not empty? Is it worth adding a comment about what should happen in this case?

llvm/test/Transforms/LoopVectorize/AArch64/scalable-call.ll

124

nit: should this test be also be called no_mapping if it is legal to vectorize?

[david-arm]

• Improved and added more comments to code + tests to explain what's going on.
• Tried to restructure the code a little to be a bit clearer.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5672

Yeah, so in the majority of cases I expect we don't have mappings for things like sqrt and sin intrinsics, at least not until there is some scalable vector library support. I've tried to restructure the code here to make the logic a bit clearer.

llvm/test/Transforms/LoopVectorize/AArch64/scalable-call.ll

124

Hi @kmclaughlin, so in this case there are no mappings attached to the call instruction, however it is legal to vectorise because there is hardware support for it. I've added comments that hopefully explains it better!

[sdesmalen]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1708

I'd rather write this the other way around, have a list here of intrinsics we do support and can safely return 'true' for, and fall back on BaseT's implementation for the 'default' case (which seems to be missing here), which is more pessimistic and only returns true if it knows the intrinsics can always be vectorized, i.e. the trivially vectorizable cases.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5652

MaxVF ?

5652

nit: s/isSclableLoopLegal/loopCanBeWidenedWithScalableVectors/ ?

5654

for eligible vectorization factors up to MaxVF

5655

Can you move this comment down to the CallInst case?

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5655

Yeah sure, although it's worth pointing out this is an existing comment that was referring to canVectorizeReductions, i.e. that we can also filter out the VFs not valid for the reductions too. I'm happy to move it down to the CallInst case though if you're not worried about the reduction case atm?

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5655

Ah I see. Looking at the implementation for canVectorizeReductions, and correspondingly to isLegalToVectorizeReduction, the VF itself isn't really used, so perhaps we may just as well remove it at this point. It's more the reduction-operation that is holding it back from vectorizing rather than the VF. Initially I imagined us having a function like you added here, which would filter the VFs based on capabilities, but it turned out that for the reductions that just wasn't really needed, but we kept the individual VF operand.

[david-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1708

So we do fall back on the BaseT version, it's just I didn't do it explicitly in the default: case. For some reason I thought this coding structure was the preferred way - I'm never really too sure about the coding style of switch statements in LLVM!

It's also worth pointing out that that the BaseT version calls isTriviallyVectorizable where even sin and cos will return true. Therefore, returning the BaseT version for those cases is not what we want. I could either:

1. Fill in the complete set of entries here with both the false and true cases, or
2. Change the BaseT version to always return false.

Any preferences?

[sdesmalen]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1708

It's also worth pointing out that that the BaseT version calls isTriviallyVectorizable where even sin and cos will return true. Therefore, returning the BaseT version for those cases is not what we want.

That suggests to me that the BaseT implementation for isLegalScalableVectorIntrinsic is too lenient and shouldn't rely on isTriviallyVectorizable. Returning 'false' seems like a better starting point.

[david-arm]

• Changed AArch64TargetTransformInfo function to return true for supported cases and default to the BaseT version for the rest.
• Changed the BaseT version of isLegalScalableVectorIntrinsic to return false.
• Removed a FIXME comment.

[sdesmalen]

llvm/include/llvm/Analysis/TargetTransformInfo.h

665

nit: I read isLegalScalableVectorIntrinsic as asking if the "scalable vector intrinsic" is legal as if the behaviour of the intrinsic is specific to scalable vectors. Rather, what you're asking is if the intrinsic (not specific to scalable vectors) is legal to use with scalable vector arguments, so from that reasoning I prefer isIntrinsicLegalForScalableVectors.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5667

When debug-info is enabled, that will insert llvm.dbg intrinsic calls which aren't handled in your switch, but maybe it's better to avoid getting into that to begin with by discarding instructions where I.isDebugOrPseudoInst() == true.

5680

Instead of passing in Msg, can you just call reportVectorizationInfo here? Also, can you make the message more specific by telling which function is not vectorizable with scalable vectors?

5682

Shouldn't the code traverse the mappings in the VFdatabase to see if there is a scalable form available?

[david-arm]

• Rebase + addressed review comments

[sdesmalen]

llvm/include/llvm/Analysis/TargetTransformInfo.h

1523

nit: Ah I only now see it falls a bit out of style, so perhaps isLegalIntrinsicForScalableVectors is a better name. Sorry I didn't spot that before.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1710

nit: maybe drop the comments, because it's not complete (for example, smax is not bit-manipulation)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1508

nit: is this an unrelated change?

5684–5691

nit: Can you use llvm::any_of() instead of the above loop?

5698

Can you add a FIXME that says this code is still too optimistic? It depends on the chosen VF whether the loop can vectorize (vis-a-vis whether a mapping is available for that specific VF). I expect we'll want to build up a set of suitable VFs (up to and including MaxVF) that are legal, so that we can later filter out VFs that should not be considered as candidates for vectorization.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1508

No, it's needed because the caller is also marked const

5684–5691

Sure. To be honest, in general I find the above loop more readable than using lambdas, which is why I tend to find it more natural to write code this way. However, I know lots of the codebase uses them. :)

[david-arm]

• Rebase.
• Addressed review comments.

[david-arm]

This issue has already been fixed.

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
   /// Return true if the target supports masked gather.
   bool isLegalMaskedGather(Type *DataType, Align Alignment) const;
 
   /// Return true if the target supports masked compress store.
   bool isLegalMaskedCompressStore(Type *DataType) const;
   /// Return true if the target supports masked expand load.
   bool isLegalMaskedExpandLoad(Type *DataType) const;
 
+  /// Return true if the target supports vectorization of the intrinsic IID for
+  /// a given ElementCount VF.
+  bool isLegalVectorIntrinsic(Intrinsic::ID IID, ElementCount VF) const;
+
   /// Return true if the target has a unified operation to calculate division
   /// and remainder. If so, the additional implicit multiplication and
   /// subtraction required to calculate a remainder from division are free. This
   /// can enable more aggressive transformations for division and remainder than
   /// would typically be allowed using throughput or size cost models.
   bool hasDivRemOp(Type *DataType, bool IsSigned) const;
 
   /// Return true if the given instruction (assumed to be a memory access
public:
   virtual bool isLegalMaskedStore(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalMaskedLoad(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalNTStore(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalNTLoad(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalMaskedScatter(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalMaskedGather(Type *DataType, Align Alignment) = 0;
   virtual bool isLegalMaskedCompressStore(Type *DataType) = 0;
   virtual bool isLegalMaskedExpandLoad(Type *DataType) = 0;
+  virtual bool isLegalVectorIntrinsic(Intrinsic::ID IID, ElementCount VF) = 0;
   virtual bool hasDivRemOp(Type *DataType, bool IsSigned) = 0;
   virtual bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) = 0;
   virtual bool prefersVectorizedAddressing() = 0;
   virtual InstructionCost getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
                                                int64_t BaseOffset,
                                                bool HasBaseReg, int64_t Scale,
                                                unsigned AddrSpace) = 0;
   virtual bool LSRWithInstrQueries() = 0;
bool isLegalMaskedGather(Type *DataType, Align Alignment) override {
     return Impl.isLegalMaskedGather(DataType, Alignment);
   }
   bool isLegalMaskedCompressStore(Type *DataType) override {
     return Impl.isLegalMaskedCompressStore(DataType);
   }
   bool isLegalMaskedExpandLoad(Type *DataType) override {
     return Impl.isLegalMaskedExpandLoad(DataType);
   }
+  bool isLegalVectorIntrinsic(Intrinsic::ID IID, ElementCount VF) override {
+    return Impl.isLegalVectorIntrinsic(IID, VF);
+  }
+
   bool hasDivRemOp(Type *DataType, bool IsSigned) override {
     return Impl.hasDivRemOp(DataType, IsSigned);
   }
   bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) override {
     return Impl.hasVolatileVariant(I, AddrSpace);
   }
   bool prefersVectorizedAddressing() override {
     return Impl.prefersVectorizedAddressing();

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   bool isLegalMaskedScatter(Type *DataType, Align Alignment) const {
     return false;
   }
 
   bool isLegalMaskedGather(Type *DataType, Align Alignment) const {
     return false;
   }
 
+  bool isLegalVectorIntrinsic(Intrinsic::ID IID, ElementCount VF) const {
+    return isTriviallyVectorizable(IID);
+  }
+
   bool isLegalMaskedCompressStore(Type *DataType) const { return false; }
 
   bool isLegalMaskedExpandLoad(Type *DataType) const { return false; }
 
   bool hasDivRemOp(Type *DataType, bool IsSigned) const { return false; }
 
   bool hasVolatileVariant(Instruction *I, unsigned AddrSpace) const {
     return false;

llvm/lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::isLegalMaskedCompressStore(Type *DataType) const {
   return TTIImpl->isLegalMaskedCompressStore(DataType);
 }
 
 bool TargetTransformInfo::isLegalMaskedExpandLoad(Type *DataType) const {
   return TTIImpl->isLegalMaskedExpandLoad(DataType);
 }
 
+bool TargetTransformInfo::isLegalVectorIntrinsic(Intrinsic::ID IID,
+                                                 ElementCount VF) const {
+  return TTIImpl->isLegalVectorIntrinsic(IID, VF);
+}
+
 bool TargetTransformInfo::hasDivRemOp(Type *DataType, bool IsSigned) const {
   return TTIImpl->hasDivRemOp(DataType, IsSigned);
 }
 
 bool TargetTransformInfo::hasVolatileVariant(Instruction *I,
                                              unsigned AddrSpace) const {
   return TTIImpl->hasVolatileVariant(I, AddrSpace);
 }

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

public:
 
   bool isLegalMaskedGather(Type *DataType, Align Alignment) const {
     return isLegalMaskedGatherScatter(DataType);
   }
   bool isLegalMaskedScatter(Type *DataType, Align Alignment) const {
     return isLegalMaskedGatherScatter(DataType);
   }
 
+  bool isLegalVectorIntrinsic(Intrinsic::ID IID, ElementCount VF) const;
+
   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/Target/AArch64/AArch64TargetTransformInfo.cpp

if (Kind == TTI::SK_Broadcast || Kind == TTI::SK_Transpose ||
     };
     std::pair<InstructionCost, MVT> LT = TLI->getTypeLegalizationCost(DL, Tp);
     if (const auto *Entry = CostTableLookup(ShuffleTbl, Kind, LT.second))
       return LT.first * Entry->Cost;
   }
 
   return BaseT::getShuffleCost(Kind, Tp, Mask, Index, SubTp);
 }
+
+bool AArch64TTIImpl::isLegalVectorIntrinsic(Intrinsic::ID IID,
+                                            ElementCount VF) const {
+  switch (IID) {
+  case Intrinsic::sin:
+  case Intrinsic::cos:
+  case Intrinsic::exp:
+  case Intrinsic::exp2:
+  case Intrinsic::log:
+  case Intrinsic::log10:
+  case Intrinsic::log2:
+  case Intrinsic::pow:
+  case Intrinsic::powi:
+  case Intrinsic::canonicalize:
+    // We can fall back on scalarization for fixed width vectors, but not for
+    // scalable vectors.
+    return !VF.isScalable();
+  }
+  return BaseT::isLegalVectorIntrinsic(IID, VF);
+}

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

bool isLegalGatherOrScatter(Value *V) {
     auto *Ty = getMemInstValueType(V);
     Align Align = getLoadStoreAlignment(V);
     return (LI && isLegalMaskedGather(Ty, Align)) ||
            (SI && isLegalMaskedScatter(Ty, Align));
   }
 
   /// Returns true if the target machine supports all of the reduction
   /// variables found for the given VF.
-  bool canVectorizeReductions(ElementCount VF) {
+  bool canVectorizeReductions(ElementCount VF) const {
     return (all_of(Legal->getReductionVars(), [&](auto &Reduction) -> bool {
       RecurrenceDescriptor RdxDesc = Reduction.second;
       return TTI.isLegalToVectorizeReduction(RdxDesc, VF);
     }));
   }
 
+  /// Returns true if all instructions in the loop are legal to vectorize using
+  /// a scalable vectorization factor VF. If the loop is illegal the function
+  /// returns an appropriate error remark in Msg.
+  bool isScalableLoopLegal(ElementCount VF, StringRef &Msg) const;
+
   /// Returns true if \p I is an instruction that will be scalarized with
   /// predication. Such instructions include conditional stores and
   /// instructions that may divide by zero.
   /// If a non-zero VF has been calculated, we check if I will be scalarized
   /// predication for that VF.
   bool
   isScalarWithPredication(Instruction *I,
                           ElementCount VF = ElementCount::getFixed(1)) const;
reportVectorizationFailure("Runtime stride check for small trip count",
         "this loop without such check by compiling with -Os/-Oz",
         "CantVersionLoopWithOptForSize", ORE, TheLoop);
     return true;
   }
 
   return false;
 }
 
+bool LoopVectorizationCostModel::isScalableLoopLegal(ElementCount VF,
+                                                     StringRef &Msg) const {
+  // Disable scalable vectorization if the loop contains unsupported reductions.
+  // Test that the loop-vectorizer can legalize all operations for this VF.
+  // FIXME: While for scalable vectors this is currently sufficient, this should
+  // be replaced by a more detailed mechanism that filters out specific VFs,
+  // instead of invalidating vectorization for a whole set of VFs based on the
+  // MaxVF.
+  if (!canVectorizeReductions(VF)) {
+    Msg = "Scalable vectorization not supported for the reduction "
+          "operations found in this loop.";
+    return false;
+  }
+
+  // Iterate through all instructions in the loop ensuring that is legal to
+  // vectorize with a scalable VF.
+  for (BasicBlock *BB : TheLoop->blocks()) {
+    for (Instruction &I : *BB) {
+      if (auto *CI = dyn_cast<CallInst>(&I)) {
+        Intrinsic::ID VecID = getVectorIntrinsicIDForCall(CI, TLI);
+        if (VecID && VFDatabase::getMappings(*CI).empty() &&
+            !TTI.isLegalVectorIntrinsic(VecID, VF)) {
+          Msg = "Scalable vectorization not supported for the call "
+                "instructions found in this loop";
+          return false;
+        }
+      }
+    }
+  }
+  return true;
+}
+
 ElementCount
 LoopVectorizationCostModel::getMaxLegalScalableVF(unsigned MaxSafeElements) {
   if (!TTI.supportsScalableVectors() && !ForceTargetSupportsScalableVectors) {
     reportVectorizationInfo(
         "Disabling scalable vectorization, because target does not "
         "support scalable vectors.",
         "ScalableVectorsUnsupported", ORE, TheLoop);
     return ElementCount::getScalable(0);
   }
 
   auto MaxScalableVF = ElementCount::getScalable(
       std::numeric_limits<ElementCount::ScalarTy>::max());
 
-  // Disable scalable vectorization if the loop contains unsupported reductions.
-  // Test that the loop-vectorizer can legalize all operations for this MaxVF.
-  // FIXME: While for scalable vectors this is currently sufficient, this should
-  // be replaced by a more detailed mechanism that filters out specific VFs,
-  // instead of invalidating vectorization for a whole set of VFs based on the
-  // MaxVF.
-  if (!canVectorizeReductions(MaxScalableVF)) {
-    reportVectorizationInfo(
-        "Scalable vectorization not supported for the reduction "
-        "operations found in this loop.",
-        "ScalableVFUnfeasible", ORE, TheLoop);
+  StringRef Msg;
+  if (!isScalableLoopLegal(MaxScalableVF, Msg)) {
+    reportVectorizationInfo(Msg, "ScalableVFUnfeasible", ORE, TheLoop);
     return ElementCount::getScalable(0);
   }
 
   if (Legal->isSafeForAnyVectorWidth())
     return MaxScalableVF;
 
   // Limit MaxScalableVF by the maximum safe dependence distance.
   Optional<unsigned> MaxVScale = TTI.getMaxVScale();

llvm/test/Transforms/LoopVectorize/AArch64/scalable-call.ll

-; RUN: opt -S -loop-vectorize -force-vector-interleave=1 -instcombine -mattr=+sve -mtriple aarch64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: opt -pass-remarks-missed=loop-vectorize -S -loop-vectorize -force-vector-interleave=1 \
+; RUN:     -instcombine -mattr=+sve -mtriple aarch64-unknown-linux-gnu < %s 2>%t | FileCheck %s
+; RUN: cat %t | FileCheck %s --check-prefix=CHECK-REMARKS
 
 define void @vec_load(i64 %N, double* nocapture %a, double* nocapture readonly %b) {
 ; CHECK-LABEL: @vec_load
 ; CHECK: vector.body:
 ; CHECK: %[[LOAD:.*]] = load <vscale x 2 x double>, <vscale x 2 x double>*
 ; CHECK: call <vscale x 2 x double> @foo_vec(<vscale x 2 x double> %[[LOAD]])
 entry:
   %cmp7 = icmp sgt i64 %N, 0
for.body:
   %iv.next = add nuw nsw i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, %N
   br i1 %exitcond, label %for.end, label %for.body, !llvm.loop !1
 
 for.end:
   ret void
 }
 
+; CHECK-REMARKS: Scalable vectorization not supported for the call instructions found in this loop
+define void @vec_sin_no_mapping(float* noalias nocapture %dst, float* noalias nocapture readonly %src, i64 %n) {
+; CHECK: @vec_sin_no_mapping
+; CHECK: call fast <2 x float> @llvm.sin.v2f32
+; CHECK-NOT: <vscale x
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i.07 = phi i64 [ %inc, %for.body ], [ 0, %entry ]
+  %arrayidx = getelementptr inbounds float, float* %src, i64 %i.07
+  %0 = load float, float* %arrayidx, align 4
+  %1 = tail call fast float @llvm.sin.f32(float %0)
+  %arrayidx1 = getelementptr inbounds float, float* %dst, i64 %i.07
+  store float %1, float* %arrayidx1, align 4
+  %inc = add nuw nsw i64 %i.07, 1
+  %exitcond.not = icmp eq i64 %inc, %n
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !1
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+
+define void @vec_sqrt_no_mapping(float* noalias nocapture %dst, float* noalias nocapture readonly %src, i64 %n) #0 {
+; CHECK: @vec_sqrt_no_mapping
+; CHECK: call fast <vscale x 2 x float> @llvm.sqrt.nxv2f32
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i.07 = phi i64 [ %inc, %for.body ], [ 0, %entry ]
+  %arrayidx = getelementptr inbounds float, float* %src, i64 %i.07
+  %0 = load float, float* %arrayidx, align 4
+  %1 = tail call fast float @llvm.sqrt.f32(float %0)
+  %arrayidx1 = getelementptr inbounds float, float* %dst, i64 %i.07
+  store float %1, float* %arrayidx1, align 4
+  %inc = add nuw nsw i64 %i.07, 1
+  %exitcond.not = icmp eq i64 %inc, %n
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !1
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+
 declare double @foo(double)
 declare i64 @bar(i64*)
 declare double @llvm.sin.f64(double)
+declare float @llvm.sin.f32(float)
+declare float @llvm.sqrt.f32(float)
 
 declare <vscale x 2 x double> @foo_vec(<vscale x 2 x double>)
 declare <vscale x 2 x i64> @bar_vec(<vscale x 2 x i64*>)
 declare <vscale x 2 x double> @sin_vec(<vscale x 2 x double>)
 
 attributes #0 = { "vector-function-abi-variant"="_ZGV_LLVM_Nxv_foo(foo_vec)" }
 attributes #1 = { "vector-function-abi-variant"="_ZGV_LLVM_Nxv_bar(bar_vec)" }
 attributes #2 = { "vector-function-abi-variant"="_ZGV_LLVM_Nxv_llvm.sin.f64(sin_vec)" }
 
 !1 = distinct !{!1, !2, !3}
 !2 = !{!"llvm.loop.vectorize.width", i32 2}
 !3 = !{!"llvm.loop.vectorize.scalable.enable", i1 true}