[LoopVectorize] Synthesize mask operands for vector variants as needed

[huntergr]

Summary

This patch allows loop vectorization with function calls in cases where masks are not required but the only available vectorized function variants are masked.

Some of the code was originally written by @paulwalker-arm

[reames]

I recommend you split this patch into the following patches:

• First, add support for calling the predicated call version when masking is not required for correctness. This requires all the plumbing for costing both versions of the call (if both are available), and let's you test those bits.
• Second, add support for predicated calls via scalar-with-predication. This requires speculation reasoning (which you don't appear to have), and is complicated enough to need testing on it's own. (This may partially exist already, but some of the code at least is incomplete since e.g. isPredicatedInst always returns false for calls.)
• Third, come back and add support for predication of calls via masking (as opposed to replication.)

The first and second step are independent, and could be done in either order.

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

1

Please autogenerate these tests for readability.

[huntergr]

Split out the functionality for synthesizing a mask when required, added a cost for mask generation.

[huntergr]

Thanks for the review. I'm tempted to add a masking equivalent to -force-target-supports-scalable-vectors=true in order to have target-independent tests, but I can add that in another patch.

[huntergr]

Ping.

@reames -- is this roughly what you expected for the case of allowing a masked variant to be used when no mask is required? I've added a cost for generating the mask (per-call for now, instead of potentially sharing it) so that we can compare costs for different VFs with and without a masked variant, but I think we would always prefer the non-masked variant for the same VF if a mask is not required.

In any case, I'm now working on the third patch.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

478 ↗

(On Diff #459684)

It would be good if the decision whether to used the masked or non-masked variant would be taken at the time of VPlan construction instead of during executing.

It would probably also be good to pass in the mask as operand to the recipe, especially if we want to support non-trivial masks in the future.

[huntergr]

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

478 ↗

(On Diff #459684)

So the problem I had with trying to decide up front was that you might have both masked and unmasked variants available, and the decision on which one to use left to the cost model -- which I think is calculated after VPlan construction.

For example, on AArch64 you might have a non-masked NEON variant and a masked SVE variant. If you know the implementation width is 128b, then the cost would be slightly higher for generating the mask for the SVE variant. If it's 256b or higher, it might be worth the extra cost due to additional parallelism.

Is there a (straightforward) way to tell VPlan that it may need to construct different recipes based on masked/non-masked variants being available? Or would this need some reworking of VPlan?

I did add the mask as an operand in the case where it is required though. If we can generate multiple recipes easily then it can be added to the operands when a dummy mask is required (and possibly shared if there are multiple calls, giving a more accurate cost).

[huntergr]

Updated to store the pointer to the vector function in the recipe rather than looking it up again during recipe execution. Forced generation of a plan per VF when there are variants available for those VFs. Added some new tests for masked vs. unmasked variants.

I'm not fond of adding the optional parameters LoopVectorizationCostModel::getVectorCallCost -- it feels like function lookup needs to be split out of it, but I'd like to get some feedback from others before doing so.

[huntergr]

Ping?

[david-arm]

Hi @huntergr, I've only partially reviewed this I'm afraid, but here are the comments I have so far. I'll try to review more this week!

llvm/include/llvm/Analysis/VectorUtils.h

129

Might be worth adding /// comments here, since the others all have them?

134

I think this will be compiled away in a release build, right? So really it's just a non-release wrapper around getParamIndexForOptionalMask. Given it's only called in one place is it worth just making getParamIndexForOptionalMask public instead and putting an assert in LoopVectorize.cpp that a mask exists?

137

This function is never called - can it be deleted?

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3442

This looks a little strange to me. In my mind, the ability to emit an active lane mask based on two integer inputs is orthogonal to how cheap it is to broadcast a true bit across a predicate. For example, an architecture may cheaply support the latter, but not the former. Maybe X86 is such an example? Can we not just let the mask cost decide the behaviour? That way you can simplify this to just

if (!VecFunc) {
   ...

3456

Do we really need both the Variant and the NeedToScalarize parameter? It looks naively that setting Variant = VecFunc is synonymous with NeedToScalarize = false. I haven't looked into this in detail so I could be wrong, but it might make more sense to remove the NeedToScalarize in favour of setting Variant?

8359

Doesn't this mean we may end up picking the least optimal VF? For example, if there are v2i32 and v4i32 masked variants we'll only ever pick the v2i32, i.e. the lowest VF?

[huntergr]

Rebased, changed as requested.

[huntergr]

llvm/include/llvm/Analysis/VectorUtils.h

137

There's now a use for it in the assert when building a recipe. (This was used in the original patch, but was left in when splitting into 3 parts).

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3442

My thinking was to treat the capability to emit an active lane mask as a proxy for being able to use masks at all, but perhaps that's a little too conservative.

I don't know if we should add a proper TTI interface to represent that capability, or just rely on the VFDatabase only having entries which the target is capable of supporting.

In any case, I've removed that check for now.

8359

No. Since we now store the pointer to the Function in the recipe, we need to force vplan to generate different plans for each VF that has a vector variant available.

See the vplan checks for 'test_v2_v4m' in synthesize-mask-for-call.ll -- there are separate VF=2 and VF=4 plans, with a widened call to different functions.

[david-arm]

Thanks a lot for addressing all the comments @huntergr! I just have a few more minor comments then I think it's good to go. :)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8359

OK, can you add some comments here explaining why you are forcing the creation of a new vplan for every subsequent VF after discovering a vector variant?

8363

nit: Could you change this to

VFShape Shape = VFShape::get(*CI, VariantVF, /*HasGlobalPred=*/true);

llvm/lib/Transforms/Vectorize/VPlan.h

949

Can you add some comments explaining what this is please? For example, that there should be one recipe for every VF because the variant requires a 1:1 mapping with the VF?

llvm/test/Transforms/LoopVectorize/AArch64/synthesize-mask-for-call.ll

8 ↗

(On Diff #490161)

I think it's also worth having a test for the case when both VF=2 and VF=4 use unmasked variants, because even in this case we must create separate vplans for each. The Variant member added to the recipe requires it now I think.

[huntergr]

Rebased, addressed latest requests.

[david-arm]

LGTM! Thanks for making all the changes @huntergr. :)

llvm/include/llvm/Analysis/VectorUtils.h

 };
 
 /// Holds the VFShape for a specific scalar to vector function mapping.
 struct VFInfo {
   VFShape Shape;          /// Classification of the vector function.
   std::string ScalarName; /// Scalar Function Name.
   std::string VectorName; /// Vector Function Name associated to this VFInfo.
   VFISAKind ISA;          /// Instruction Set Architecture.
+
+  unsigned getParamIndexForMask() const {
+    auto MaskPos = getParamIndexForOptionalMask();
+    if (MaskPos)
+      return *MaskPos;
+
+    llvm_unreachable("Requested paramater index of non-existent mask!");
+  }
+
+  bool isMasked() const { return getParamIndexForOptionalMask().has_value(); }
+
+private:
+  Optional<unsigned> getParamIndexForOptionalMask() const {
+    unsigned ParamCount = Shape.Parameters.size();
+    for (unsigned i = 0; i < ParamCount; ++i)
+      if (Shape.Parameters[i].ParamKind == VFParamKind::GlobalPredicate)
+        return i;
+
+    return None;
+  }
 };
 
 namespace VFABI {
 /// LLVM Internal VFABI ISA token for vector functions.
 static constexpr char const *_LLVM_ = "_LLVM_";
 /// Prefix for internal name redirection for vector function that
 /// tells the compiler to scalarize the call using the scalar name
 /// of the function. For example, a mangled name like

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

for (Instruction &I : *BB) {
     }
 
     // Do not let llvm.experimental.noalias.scope.decl block the vectorization.
     // TODO: there might be cases that it should block the vectorization. Let's
     // ignore those for now.
     if (isa<NoAliasScopeDeclInst>(&I))
       continue;
 
+    if (CallInst *CI = dyn_cast<CallInst>(&I)) {
+      // Check whether we have at least one masked vector version of a scalar
+      // function.
+      bool HasMaskedVersion = false;
+
+      auto Mappings = VFDatabase::getMappings(*CI);
+      for (VFInfo Info : Mappings)
+        HasMaskedVersion |= Info.isMasked();
+
+      if (HasMaskedVersion) {
+        MaskedOp.insert(CI);
+        continue;
+      }
+    }
+
     // Loads are handled via masking (or speculated if safe to do so.)
     if (auto *LI = dyn_cast<LoadInst>(&I)) {
       if (!SafePtrs.count(LI->getPointerOperand()))
         MaskedOp.insert(LI);
       continue;
     }
 
     // Predicated store requires some form of masking:

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// loop and the start value for the canonical induction, if it is != 0. The
   /// latter is the case when vectorizing the epilogue loop. In the case of
   /// epilogue vectorization, this function is overriden to handle the more
   /// complex control flow around the loops.
   virtual std::pair<BasicBlock *, Value *> createVectorizedLoopSkeleton();
 
   /// Widen a single call instruction within the innermost loop.
   void widenCallInstruction(CallInst &CI, VPValue *Def, VPUser &ArgOperands,
-                            VPTransformState &State);
+                            VPTransformState &State, bool MaskAvailable);
 
   /// Fix the vectorized code, taking care of header phi's, live-outs, and more.
   void fixVectorizedLoop(VPTransformState &State, VPlan &Plan);
 
   // Return true if any runtime check is added.
   bool areSafetyChecksAdded() { return AddedSafetyChecks; }
 
   /// A type for vectorized values in the new loop. Each value from the
public:
   InstructionCost getVectorIntrinsicCost(CallInst *CI, ElementCount VF) const;
 
   /// Estimate cost of a call instruction CI if it were vectorized with factor
   /// VF. Return the cost of the instruction, including scalarization overhead
   /// if it's needed. The flag NeedToScalarize shows if the call needs to be
   /// scalarized -
   /// i.e. either vector version isn't available, or is too expensive.
   InstructionCost getVectorCallCost(CallInst *CI, ElementCount VF,
+                                    bool NeedsMask,
                                     bool &NeedToScalarize) const;
 
   /// Returns true if the per-lane cost of VectorizationFactor A is lower than
   /// that of B.
   bool isMoreProfitable(const VectorizationFactor &A,
                         const VectorizationFactor &B) const;
 
   /// Invalidates decisions already taken by the cost model.
for (Instruction &In : llvm::make_early_inc_range(*BB)) {
     }
 
     CSEMap[&In] = &In;
   }
 }
 
 InstructionCost
 LoopVectorizationCostModel::getVectorCallCost(CallInst *CI, ElementCount VF,
+                                              bool NeedsMask,
                                               bool &NeedToScalarize) const {
   Function *F = CI->getCalledFunction();
   Type *ScalarRetTy = CI->getType();
   SmallVector<Type *, 4> Tys, ScalarTys;
   for (auto &ArgOp : CI->args())
     ScalarTys.push_back(ArgOp->getType());
 
   // Estimate cost of scalarized vector call. The source operands are assumed
LoopVectorizationCostModel::getVectorCallCost(CallInst *CI, ElementCount VF,
   InstructionCost ScalarizationCost = getScalarizationOverhead(CI, VF);
 
   InstructionCost Cost =
       ScalarCallCost * VF.getKnownMinValue() + ScalarizationCost;
 
   // If we can't emit a vector call for this function, then the currently found
   // cost is the cost we need to return.
   NeedToScalarize = true;
-  VFShape Shape = VFShape::get(*CI, VF, false /*HasGlobalPred*/);
+  VFShape Shape = VFShape::get(*CI, VF, NeedsMask);
   Function *VecFunc = VFDatabase(*CI).getVectorizedFunction(Shape);
+  // If we want an unmasked vector function but can't find one matching the VF,
+  // and the target supports an active lane mask, maybe we can find vector
+  // function that does use a mask and synthesize an all-true mask.
+  if (!VecFunc && !NeedsMask &&
+      TTI.emitGetActiveLaneMask() != PredicationStyle::None) {
+    Shape = VFShape::get(*CI, VF, /*HasGlobalPred=*/true);
+    VecFunc = VFDatabase(*CI).getVectorizedFunction(Shape);
+  }
 
   if (!TLI || CI->isNoBuiltin() || !VecFunc)
     return Cost;
 
   // If the corresponding vector cost is cheaper, return its cost.
   InstructionCost VectorCallCost =
       TTI.getCallInstrCost(nullptr, RetTy, Tys, TTI::TCK_RecipThroughput);
   if (VectorCallCost < Cost) {
     NeedToScalarize = false;
     Cost = VectorCallCost;
   }
   return Cost;
 }
 
 static Type *MaybeVectorizeType(Type *Elt, ElementCount VF) {
   if (VF.isScalar() || (!Elt->isIntOrPtrTy() && !Elt->isFloatingPointTy()))
     return Elt;
   return VectorType::get(Elt, VF);

 
 bool InnerLoopVectorizer::useOrderedReductions(
     const RecurrenceDescriptor &RdxDesc) {
   return Cost->useOrderedReductions(RdxDesc);
 }
 
 void InnerLoopVectorizer::widenCallInstruction(CallInst &CI, VPValue *Def,
                                                VPUser &ArgOperands,
-                                               VPTransformState &State) {
+                                               VPTransformState &State,
+                                               bool MaskAvailable) {
   assert(!isa<DbgInfoIntrinsic>(CI) &&
          "DbgInfoIntrinsic should have been dropped during VPlan construction");
   State.setDebugLocFromInst(&CI);
 
-  SmallVector<Type *, 4> Tys;
-  for (Value *ArgOperand : CI.args())
-    Tys.push_back(ToVectorTy(ArgOperand->getType(), VF.getKnownMinValue()));
-
   Intrinsic::ID ID = getVectorIntrinsicIDForCall(&CI, TLI);
 
   // The flag shows whether we use Intrinsic or a usual Call for vectorized
   // version of the instruction.
   // Is it beneficial to perform intrinsic call compared to lib call?
   bool NeedToScalarize = false;
-  InstructionCost CallCost = Cost->getVectorCallCost(&CI, VF, NeedToScalarize);
+  InstructionCost CallCost =
+      Cost->getVectorCallCost(&CI, VF, MaskAvailable, NeedToScalarize);
   InstructionCost IntrinsicCost =
       ID ? Cost->getVectorIntrinsicCost(&CI, VF) : 0;
   bool UseVectorIntrinsic = ID && IntrinsicCost <= CallCost;
   assert((UseVectorIntrinsic || !NeedToScalarize) &&
          "Instruction should be scalarized elsewhere.");
   assert((IntrinsicCost.isValid() || CallCost.isValid()) &&
          "Either the intrinsic cost or vector call cost must be valid");
 
+  // If we added a mask operand in the recipe, extract it so that we can
+  // insert it in the right position for the vectorized call. The mask isn't
+  // guaranteed to be the last argument.
+  VPValue *VPMask = nullptr;
+  if (MaskAvailable)
+    VPMask = ArgOperands.removeAndReturnLastOperand();
+
   for (unsigned Part = 0; Part < UF; ++Part) {
     SmallVector<Type *, 2> TysForDecl = {CI.getType()};
     SmallVector<Value *, 4> Args;
     for (const auto &I : enumerate(ArgOperands.operands())) {
       // Some intrinsics have a scalar argument - don't replace it with a
       // vector.
       Value *Arg;
       if (!UseVectorIntrinsic ||
           !isVectorIntrinsicWithScalarOpAtArg(ID, I.index()))
         Arg = State.get(I.value(), Part);
       else
         Arg = State.get(I.value(), VPIteration(0, 0));
       if (isVectorIntrinsicWithOverloadTypeAtArg(ID, I.index()))
         TysForDecl.push_back(Arg->getType());
       Args.push_back(Arg);
     }
 
     Function *VectorF;
+    bool VectorFTakesMask = false;
+    unsigned VectorFMaskPos = 0;
+
     if (UseVectorIntrinsic) {
       // Use vector version of the intrinsic.
       if (VF.isVector())
         TysForDecl[0] = VectorType::get(CI.getType()->getScalarType(), VF);
       Module *M = State.Builder.GetInsertBlock()->getModule();
       VectorF = Intrinsic::getDeclaration(M, ID, TysForDecl);
       assert(VectorF && "Can't retrieve vector intrinsic.");
     } else {
       // Use vector version of the function call.
-      const VFShape Shape = VFShape::get(CI, VF, false /*HasGlobalPred*/);
+      VFShape Shape = VFShape::get(CI, VF, MaskAvailable);
+
+      VectorF = VFDatabase(CI).getVectorizedFunction(Shape);
+
+      if (!VectorF && !MaskAvailable &&
+          TTI->emitGetActiveLaneMask() != PredicationStyle::None) {
+        Shape = VFShape::get(CI, VF, /*HasGlobalPred=*/true);
+        VectorF = VFDatabase(CI).getVectorizedFunction(Shape);
+      }
 #ifndef NDEBUG
-      assert(VFDatabase(CI).getVectorizedFunction(Shape) != nullptr &&
-             "Can't create vector function.");
+      assert(VectorF != nullptr && "Can't create vector function.");
 #endif
-      VectorF = VFDatabase(CI).getVectorizedFunction(Shape);
+      // Check the VFInfo for masking details
+      for (VFInfo Info : VFDatabase(CI).getMappings(CI)) {
+        if (Info.Shape == Shape) {
+            VectorFTakesMask = Info.isMasked();
+            if (VectorFTakesMask)
+            VectorFMaskPos = Info.getParamIndexForMask();
+            break;
+        }
+      }
+    }
+
+    assert((!MaskAvailable || VectorFTakesMask) &&
+           "Mask supplied for function with no mask argument");
+
+    if (VectorFTakesMask) {
+      Value *Mask = nullptr;
+      if (VPMask)
+        Mask = State.get(VPMask, Part);
+      else
+        Mask = ConstantInt::getTrue(VectorType::get(
+            IntegerType::getInt1Ty(VectorF->getFunctionType()->getContext()),
+            VF));
+      Args.insert(Args.begin() + VectorFMaskPos, Mask);
     }
-      SmallVector<OperandBundleDef, 1> OpBundles;
-      CI.getOperandBundlesAsDefs(OpBundles);
-      CallInst *V = Builder.CreateCall(VectorF, Args, OpBundles);
+
 
-      if (isa<FPMathOperator>(V))
-        V->copyFastMathFlags(&CI);
 
-      State.set(Def, V, Part);
-      State.addMetadata(V, &CI);
   }
 }
 
 void LoopVectorizationCostModel::collectLoopScalars(ElementCount VF) {
   // We should not collect Scalars more than once per VF. Right now, this
   // function is called from collectUniformsAndScalars(), which already does
   // this check. Collecting Scalars for VF=1 does not make any sense.
   assert(VF.isVector() && Scalars.find(VF) == Scalars.end() &&
case Instruction::FPTrunc: {
     return TTI.getCastInstrCost(Opcode, VectorTy, SrcVecTy, CCH, CostKind, I);
   }
   case Instruction::Call: {
     if (RecurrenceDescriptor::isFMulAddIntrinsic(I))
       if (auto RedCost = getReductionPatternCost(I, VF, VectorTy, CostKind))
         return *RedCost;
     bool NeedToScalarize;
     CallInst *CI = cast<CallInst>(I);
-    InstructionCost CallCost = getVectorCallCost(CI, VF, NeedToScalarize);
+    InstructionCost CallCost =
+        getVectorCallCost(CI, VF, Legal->isMaskRequired(CI), NeedToScalarize);
     if (getVectorIntrinsicIDForCall(CI, TLI)) {
       InstructionCost IntrinsicCost = getVectorIntrinsicCost(CI, VF);
       return std::min(CallCost, IntrinsicCost);
     }
     return CallCost;
   }
   case Instruction::ExtractValue:
     return TTI.getInstructionCost(I, TTI::TCK_RecipThroughput);
for (unsigned In = 0; In < NumIncoming; In++) {
     if (EdgeMask)
       OperandsWithMask.push_back(EdgeMask);
   }
   return toVPRecipeResult(new VPBlendRecipe(Phi, OperandsWithMask));
 }
 
 VPWidenCallRecipe *VPRecipeBuilder::tryToWidenCall(CallInst *CI,
                                                    ArrayRef<VPValue *> Operands,
-                                                   VFRange &Range) const {
+                                                   VFRange &Range,
+                                                   VPlanPtr &Plan) {
 
   bool IsPredicated = LoopVectorizationPlanner::getDecisionAndClampRange(
       [this, CI](ElementCount VF) {
         return CM.isScalarWithPredication(CI, VF);
       },
       Range);
 
   if (IsPredicated)
     return nullptr;
 
+  VPValue *Mask = nullptr;
+  if (Legal->isMaskRequired(CI))
+    Mask = createBlockInMask(CI->getParent(), Plan);
+
   Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
   if (ID && (ID == Intrinsic::assume || ID == Intrinsic::lifetime_end ||
              ID == Intrinsic::lifetime_start || ID == Intrinsic::sideeffect ||
              ID == Intrinsic::pseudoprobe ||
              ID == Intrinsic::experimental_noalias_scope_decl))
     return nullptr;
 
   auto willWiden = [&](ElementCount VF) -> bool {
     Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
     // The following case may be scalarized depending on the VF.
     // The flag shows whether we use Intrinsic or a usual Call for vectorized
     // version of the instruction.
     // Is it beneficial to perform intrinsic call compared to lib call?
     bool NeedToScalarize = false;
-    InstructionCost CallCost = CM.getVectorCallCost(CI, VF, NeedToScalarize);
+    InstructionCost CallCost = CM.getVectorCallCost(
+        CI, VF, Legal->isMaskRequired(CI), NeedToScalarize);
     InstructionCost IntrinsicCost = ID ? CM.getVectorIntrinsicCost(CI, VF) : 0;
     bool UseVectorIntrinsic = ID && IntrinsicCost <= CallCost;
     return UseVectorIntrinsic || !NeedToScalarize;
   };
 
   if (!LoopVectorizationPlanner::getDecisionAndClampRange(willWiden, Range))
     return nullptr;
 
   ArrayRef<VPValue *> Ops = Operands.take_front(CI->arg_size());
-  return new VPWidenCallRecipe(*CI, make_range(Ops.begin(), Ops.end()));
+  return new VPWidenCallRecipe(*CI, make_range(Ops.begin(), Ops.end()), Mask);
 }
 
 bool VPRecipeBuilder::shouldWiden(Instruction *I, VFRange &Range) const {
   assert(!isa<BranchInst>(I) && !isa<PHINode>(I) && !isa<LoadInst>(I) &&
          !isa<StoreInst>(I) && "Instruction should have been handled earlier");
   // Instruction should be widened, unless it is scalar after vectorization,
   // scalarization is profitable or it is predicated.
   auto WillScalarize = [this, I](ElementCount VF) -> bool {
if (isa<TruncInst>(Instr) &&
     return toVPRecipeResult(Recipe);
 
   // All widen recipes below deal only with VF > 1.
   if (LoopVectorizationPlanner::getDecisionAndClampRange(
           [&](ElementCount VF) { return VF.isScalar(); }, Range))
     return nullptr;
 
   if (auto *CI = dyn_cast<CallInst>(Instr))
-    return toVPRecipeResult(tryToWidenCall(CI, Operands, Range));
+    return toVPRecipeResult(tryToWidenCall(CI, Operands, Range, Plan));
 
   if (isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
     return toVPRecipeResult(tryToWidenMemory(Instr, Operands, Range, Plan));
 
   if (!shouldWiden(Instr, Range))
     return nullptr;
 
   if (auto GEP = dyn_cast<GetElementPtrInst>(Instr))
for (unsigned i = 0; i < IG->getFactor(); ++i) {
     }
     ++OpIdx;
   }
 }
 #endif
 
 void VPWidenCallRecipe::execute(VPTransformState &State) {
   State.ILV->widenCallInstruction(*cast<CallInst>(getUnderlyingInstr()), this,
-                                  *this, State);
+                                  *this, State, Mask);
 }
 
 void VPWidenIntOrFpInductionRecipe::execute(VPTransformState &State) {
   assert(!State.Instance && "Int or FP induction being replicated.");
 
   Value *Start = getStartValue()->getLiveInIRValue();
   const InductionDescriptor &ID = getInductionDescriptor();
   TruncInst *Trunc = getTruncInst();

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

class VPRecipeBuilder {
   /// performs full if-conversion.
   VPRecipeOrVPValueTy tryToBlend(PHINode *Phi, ArrayRef<VPValue *> Operands,
                                  VPlanPtr &Plan);
 
   /// Handle call instructions. If \p CI can be widened for \p Range.Start,
   /// return a new VPWidenCallRecipe. Range.End may be decreased to ensure same
   /// decision from \p Range.Start to \p Range.End.
   VPWidenCallRecipe *tryToWidenCall(CallInst *CI, ArrayRef<VPValue *> Operands,
-                                    VFRange &Range) const;
+                                    VFRange &Range, VPlanPtr &Plan);
 
   /// Check if \p I has an opcode that can be widened and return a VPWidenRecipe
   /// if it can. The function should only be called if the cost-model indicates
   /// that widening should be performed.
   VPWidenRecipe *tryToWiden(Instruction *I, ArrayRef<VPValue *> Operands) const;
 
   /// Return a VPRecipeOrValueTy with VPRecipeBase * being set. This can be used to force the use as VPRecipeBase* for recipe sub-types that also inherit from VPValue.
   VPRecipeOrVPValueTy toVPRecipeResult(VPRecipeBase *R) const { return R; }

llvm/lib/Transforms/Vectorize/VPlan.h

#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 #endif
 };
 
 /// A recipe for widening Call instructions.
 class VPWidenCallRecipe : public VPRecipeBase, public VPValue {
+  bool Mask;
 
 public:
   template <typename IterT>
-  VPWidenCallRecipe(CallInst &I, iterator_range<IterT> CallArguments)
+  VPWidenCallRecipe(CallInst &I, iterator_range<IterT> CallArguments,
+                    VPValue *MaskVal = nullptr)
       : VPRecipeBase(VPRecipeBase::VPWidenCallSC, CallArguments),
-        VPValue(VPValue::VPVWidenCallSC, &I, this) {}
+        VPValue(VPValue::VPVWidenCallSC, &I, this), Mask(MaskVal != nullptr) {
+    if (MaskVal)
+      addOperand(MaskVal);
+  }
 
   ~VPWidenCallRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     return D->getVPDefID() == VPRecipeBase::VPWidenCallSC;
   }
 

llvm/lib/Transforms/Vectorize/VPlanValue.h

void setOperand(unsigned I, VPValue *New) {
     New->addUser(*this);
   }
 
   void removeLastOperand() {
     VPValue *Op = Operands.pop_back_val();
     Op->removeUser(*this);
   }
 
+  VPValue *removeAndReturnLastOperand() {
+    VPValue *Op = Operands.pop_back_val();
+    Op->removeUser(*this);
+    return Op;
+  }
+
   typedef SmallVectorImpl<VPValue *>::iterator operand_iterator;
   typedef SmallVectorImpl<VPValue *>::const_iterator const_operand_iterator;
   typedef iterator_range<operand_iterator> operand_range;
   typedef iterator_range<const_operand_iterator> const_operand_range;
 
   operand_iterator op_begin() { return Operands.begin(); }
   const_operand_iterator op_begin() const { return Operands.begin(); }
   operand_iterator op_end() { return Operands.end(); }

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

 ; RUN: opt < %s -loop-vectorize -force-vector-interleave=1 -instsimplify -S | FileCheck %s --check-prefixes=CHECK,LV
 ; RUN: opt < %s -loop-vectorize -force-vector-interleave=1 -prefer-predicate-over-epilogue=predicate-dont-vectorize -instsimplify -S | FileCheck %s --check-prefixes=CHECK,TFALWAYS
 ; RUN: opt < %s -loop-vectorize -force-vector-interleave=1 -prefer-predicate-over-epilogue=predicate-else-scalar-epilogue -instsimplify -S | FileCheck %s --check-prefixes=CHECK,TFFALLBACK
 
 target triple = "aarch64-unknown-linux-gnu"
 
 ; A call whose argument must be widened. We check that tail folding uses the
 ; primary mask, and that without tail folding we synthesize an all-true mask.
 define void @test_widen(i64* noalias %a, i64* readnone %b) #4 {
 ; CHECK-LABEL: @test_widen(
-; LV-NOT: call <vscale x 2 x i64> @foo_vector
-; TFALWAYS-NOT: vector.body
-; TFALWAYS-NOT: call <vscale x 2 x i64> @foo_vector
-; TFFALLBACK-NOT: call <vscale x 2 x i64> @foo_vector
+; LV: %[[LOAD:.+]] = load <vscale x 2 x i64>
+; LV: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %[[LOAD]], <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer))
+; TFALWAYS: %[[MASK:.+]] = phi <vscale x 2 x i1>
+; TFALWAYS: %[[LOAD:.+]] = call <vscale x 2 x i64> @llvm.masked.load.nxv2i64.p0nxv2i64
+; TFALWAYS: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %[[LOAD]], <vscale x 2 x i1> %[[MASK]])
+; TFFALLBACK: %[[MASK:.+]] = phi <vscale x 2 x i1>
+; TFFALLBACK: %[[LOAD:.+]] = call <vscale x 2 x i64> @llvm.masked.load.nxv2i64.p0nxv2i64
+; TFFALLBACK: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %[[LOAD]], <vscale x 2 x i1> %[[MASK]])
 ; CHECK: ret void
 entry:
   br label %for.body
 
 for.body:
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %gep = getelementptr i64, i64* %b, i64 %indvars.iv
   %load = load i64, i64* %gep
   %call = call i64 @foo(i64 %load) #1
   %arrayidx = getelementptr inbounds i64, i64* %a, i64 %indvars.iv
   store i64 %call, i64* %arrayidx
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv.next, 1024
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 
 for.cond.cleanup:
   ret void
 }
 
 ; Check that a simple conditional call can be vectorized.
 define void @test_if_then(i64* noalias %a, i64* readnone %b) #4 {
 ; CHECK-LABEL: @test_if_then(
-; LV-NOT: call <vscale x 2 x i64> @foo_vector
-; TFALWAYS-NOT: call <vscale x 2 x i64> @foo_vector
-; TFFALLBACK-NOT: call <vscale x 2 x i64> @foo_vector
+; LV: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %wide.load, <vscale x 2 x i1> %{{.+}})
+; TFALWAYS: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %wide.masked.load, <vscale x 2 x i1> %{{.+}})
+; TFFALLBACK: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %wide.masked.load, <vscale x 2 x i1> %{{.+}})
 ; CHECK: ret void
 entry:
   br label %for.body
 
 for.body:
   %indvars.iv = phi i64 [ %indvars.iv.next, %if.end ], [ 0, %entry ]
   %arrayidx = getelementptr inbounds i64, i64* %a, i64 %indvars.iv
   %0 = load i64, i64* %arrayidx, align 8

 }
 
 ; This checks the ability to handle masking of an if-then-else CFG with
 ; calls inside the conditional blocks. Although one of the calls has a
 ; uniform parameter and the metadata lists a uniform variant, right now
 ; we just see a splat of the parameter instead. More work needed.
 define void @test_widen_if_then_else(i64* noalias %a, i64* readnone %b) #4 {
 ; CHECK-LABEL: @test_widen_if_then_else
-; LV-NOT: call <vscale x 2 x i64> @foo_vector
-; LV-NOT: call <vscale x 2 x i64> @foo_uniform
-; TFALWAYS-NOT: call <vscale x 2 x i64> @foo_vector
-; TFALWAYS-NOT: call <vscale x 2 x i64> @foo_uniform
-; TFFALLBACK-NOT: call <vscale x 2 x i64> @foo_vector
-; TFFALLBACK-NOT: call <vscale x 2 x i64> @foo_uniform
+; LV: %[[LOAD:.+]] = load <vscale x 2 x i64>, <vscale x 2 x i64>* %{{[0-9]+}}
+; LV: %[[CMP:.+]] = icmp ugt <vscale x 2 x i64> %[[LOAD]], shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 50, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
+; LV: %[[INV:.+]] = xor <vscale x 2 x i1> %[[CMP]], shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer)
+; LV: %[[UNIFORM_SPLAT:.+]] = call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> zeroinitializer, <vscale x 2 x i1> %[[INV]])
+; LV: %[[VECTOR:.+]] = call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %[[LOAD]], <vscale x 2 x i1> %[[CMP]])
+; LV: %[[PPHI:.+]] = select <vscale x 2 x i1> %[[INV]], <vscale x 2 x i64> %[[UNIFORM_SPLAT]], <vscale x 2 x i64> %[[VECTOR]]
+; LV: store <vscale x 2 x i64> %[[PPHI]]
+; TFALWAYS: %[[LOAD:.+]] = call <vscale x 2 x i64> @llvm.masked.load.nxv2i64.p0nxv2i64(<vscale x 2 x i64>* %{{.+}}, i32 8, <vscale x 2 x i1> %active.lane.mask, <vscale x 2 x i64> poison)
+; TFALWAYS: %[[CMP:.+]] = icmp ugt <vscale x 2 x i64> %[[LOAD]], shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 50, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
+; TFALWAYS: %[[INV:.+]] = xor <vscale x 2 x i1> %[[CMP]], shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer)
+; TFALWAYS: %[[MERGE1:.+]] = select <vscale x 2 x i1> %active.lane.mask, <vscale x 2 x i1> %[[INV]], <vscale x 2 x i1> zeroinitializer
+; TFALWAYS: %[[UNIFORM_SPLAT:.+]] = call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> zeroinitializer, <vscale x 2 x i1> %[[MERGE1]])
+; TFALWAYS: %[[MERGE2:.+]] = select <vscale x 2 x i1> %active.lane.mask, <vscale x 2 x i1> %[[CMP]], <vscale x 2 x i1> zeroinitializer
+; TFALWAYS: %[[VECTOR:.+]] = call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %[[LOAD]], <vscale x 2 x i1> %[[MERGE2]])
+; TFALWAYS: %[[PPHI:.+]] = select <vscale x 2 x i1> %[[MERGE1]], <vscale x 2 x i64> %[[UNIFORM_SPLAT]], <vscale x 2 x i64> %[[VECTOR]]
+; TFALWAYS: call void @llvm.masked.store.nxv2i64.p0nxv2i64(<vscale x 2 x i64> %[[PPHI]]
+; TFFALLBACK: %[[LOAD:.+]] = call <vscale x 2 x i64> @llvm.masked.load.nxv2i64.p0nxv2i64(<vscale x 2 x i64>* %{{.+}}, i32 8, <vscale x 2 x i1> %active.lane.mask, <vscale x 2 x i64> poison)
+; TFFALLBACK: %[[CMP:.+]] = icmp ugt <vscale x 2 x i64> %[[LOAD]], shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 50, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
+; TFFALLBACK: %[[INV:.+]] = xor <vscale x 2 x i1> %[[CMP]], shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer)
+; TFFALLBACK: %[[MERGE1:.+]] = select <vscale x 2 x i1> %active.lane.mask, <vscale x 2 x i1> %[[INV]], <vscale x 2 x i1> zeroinitializer
+; TFFALLBACK: %[[UNIFORM_SPLAT:.+]] = call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> zeroinitializer, <vscale x 2 x i1> %[[MERGE1]])
+; TFFALLBACK: %[[MERGE2:.+]] = select <vscale x 2 x i1> %active.lane.mask, <vscale x 2 x i1> %[[CMP]], <vscale x 2 x i1> zeroinitializer
+; TFFALLBACK: %[[VECTOR:.+]] = call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %[[LOAD]], <vscale x 2 x i1> %[[MERGE2]])
+; TFFALLBACK: %[[PPHI:.+]] = select <vscale x 2 x i1> %[[MERGE1]], <vscale x 2 x i64> %[[UNIFORM_SPLAT]], <vscale x 2 x i64> %[[VECTOR]]
+; TFFALLBACK: call void @llvm.masked.store.nxv2i64.p0nxv2i64(<vscale x 2 x i64> %[[PPHI]]
 ; CHECK: ret void
 entry:
   br label %for.body
 
 for.body:
   %indvars.iv = phi i64 [ %indvars.iv.next, %if.end ], [ 0, %entry ]
   %arrayidx = getelementptr inbounds i64, i64* %a, i64 %indvars.iv
   %0 = load i64, i64* %arrayidx, align 8
for.cond.cleanup:
   ret void
 }
 
 ; A call whose argument must be widened, where the vector variant does not have
 ; a mask. Forcing tail folding results in no vectorized call, whereas an
 ; unpredicated body with scalar tail can use the unmasked variant.
 define void @test_widen_nomask(i64* noalias %a, i64* readnone %b) #4 {
 ; CHECK-LABEL: @test_widen_nomask(
-; LV: call <vscale x 2 x i64> @foo_vector_nomask
-; TFALWAYS-NOT: vector.body
+; LV: call <vscale x 2 x i64> @foo_vector_nomask(<vscale x 2 x i64> %wide.load)
 ; TFALWAYS-NOT: call <vscale x 2 x i64> @foo_vector_nomask
-; TFFALLBACK: call <vscale x 2 x i64> @foo_vector_nomask
+; TFFALLBACK: call <vscale x 2 x i64> @foo_vector_nomask(<vscale x 2 x i64> %wide.load)
 ; CHECK: ret void
 entry:
   br label %for.body
 
 for.body:
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %gep = getelementptr i64, i64* %b, i64 %indvars.iv
   %load = load i64, i64* %gep
   %call = call i64 @foo(i64 %load) #2
   %arrayidx = getelementptr inbounds i64, i64* %a, i64 %indvars.iv
   store i64 %call, i64* %arrayidx
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv.next, 1024
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 
 for.cond.cleanup:
   ret void
 }
 
 ; If both masked and unmasked options are present, we expect to see tail folding
 ; use the masked version and unpredicated body with scalar tail use the unmasked
 ; version.
 define void @test_widen_optmask(i64* noalias %a, i64* readnone %b) #4 {
 ; CHECK-LABEL: @test_widen_optmask(
-; LV: call <vscale x 2 x i64> @foo_vector_nomask
-; TFALWAYS-NOT: vector.body
-; TFALWAYS-NOT: call <vscale x 2 x i64> @foo_vector
-; TFFALLBACK: call <vscale x 2 x i64> @foo_vector_nomask
+; LV: call <vscale x 2 x i64> @foo_vector_nomask(<vscale x 2 x i64> %wide.load)
+; TFALWAYS: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %wide.masked.load, <vscale x 2 x i1> %active.lane.mask)
+; TFFALLBACK: call <vscale x 2 x i64> @foo_vector(<vscale x 2 x i64> %wide.masked.load, <vscale x 2 x i1> %active.lane.mask)
 ; CHECK: ret void
 entry:
   br label %for.body
 
 for.body:
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %gep = getelementptr i64, i64* %b, i64 %indvars.iv
   %load = load i64, i64* %gep