[SVE][LoopVectorize] Add support for extracting the last lane of a scalable vector

[david-arm]

Summary

There are certain loops like this below:

for (int i = 0; i < n; i++) {
  a[i] = b[i] + 1;
  *inv = a[i];
}

that can only be vectorised if we are able to extract the last lane of the
vectorised form of 'a[i]'. For fixed width vectors this already works since
we know at compile time what the final lane is, however for scalable vectors
this is a different story. This patch adds support for extracting the last
lane from a scalable vector using a runtime determined lane value. I have
added support to VPIteration for non-constant lanes that still permits the
caching of values. Whilst doing this work I couldn't find any explicit tests
for extracting the last lane values of fixed width vectors so I added tests
for both scalable and fixed width vectors.

[sdesmalen]

llvm/lib/Transforms/Vectorize/VPlan.h

168

For scalable vectors, it probably only ever makes sense to capture any of the following lanes:

• The first N lanes from <vscale x N x <eltty>>
• The last N lanes from <vscale x N x <eltty>>

I'm not sure if the loop-vectorizer would currently ever need more than just the first/last lane, but I could imagine for interleaving it may want to extract the second/third/fourth-last value from the vector.

Perhaps you can represent this with:

unsigned LaneIdx;
enum {
  LK_Fixed,
  LK_ScalableFirst,
  LK_ScalableLast,
} LaneKind;

?

[david-arm]

llvm/lib/Transforms/Vectorize/VPlan.h

168

I'm happy with the idea of adding an extra member to VPInstance that contains an enum and is probably nicer than what I have now! I'm not sure if we need a LK_ScalableFirst though as this is always known at compile time to be 0 I think - perhaps we just need a LK_First and a LK_ScalableLast? Also, are suggesting that the enum describes how to use Lane, i.e.

if (StartFromFirst)

Index = Lane

else if (StartFromLast)

Index = NumElts - 1 - Lane?

or with LK_ScalableLast do you literally mean the last lane of the vector?

[david-arm]

• Following @sdesmalen's suggestion I've added an enum to VPIteration that describes how we should interpret the 'Lane' member variable. This means we can support offsets from the last lane up to the known minimum number of elements.

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2486

should this be implied by Instance.Lane > 0? (implemented with operator overload, Instance > 0)

2507–2513

nit:

switch (Instance.Kind) {
case VPIteration::LK_First:
  Lane = Builder.getInt32(Instance.Lane);
  break;
case VPIteration::LK_ScalableLast:
  Lane = Builder.CreateSub(...);
  break;
}

(without default). That will give a compile-time warning when a new kind is added.

4453

Should Lane not be set to VF.getKnownMinValue() - 1 for VF.isScalable() as well?

llvm/lib/Transforms/Vectorize/VPlan.h

166–173

How about:

/// LaneKind describes how to interpret Lane.
/// For LK_First, Lane is the index into the first N elements of a fixed-vector <N x <etltty>> or a scalable vector <vscale x N x <eltty>>.
/// For LK_ScalableLast, Lane is the index into the last N elements of a scalable vector <vscale x N x  <eltty>>

168

I indeed meant that LaneKind indexes the first or last 'chunk' in the vector, e.g. v0, v1 for the first, and vN-2, vN-1 for the last in:

<vscale x 2 x i32> <=> <elt0, elt1 |  elt2, elt3 | ... | eltN-2, eltN-1>

234

This only needs to be 2 * VF.getKnownMinValue() if VF is scalable.

241

use LK_First and drop the default, so that you get a compile-time warning if a new kind is added to the enum.

243–244

Can you make this a method to VPIteration, e.g. getIndex()

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4453

Not in the way I've defined it, i.e. LastLane - Lane. This is a bit similar to how the offset is defined for Cullen's vector.splice intrinsic. A Lane of 0 is the same as LastLane - this does make the calculation of the runtime lane a bit easier.

llvm/lib/Transforms/Vectorize/VPlan.h

166–173

Sure, although I think the way I've defined LK_ScalableLast is not quite the same as you described above. The calculation of lane in my patch is:

LaneFromStartOfVec = LastLaneOfVec - Lane.

243–244

Sure. I did originally think about that, but then I wondered if that only really makes sense in the context of a cache? For example, if I move this to VPIteration then I think I should also move getNumCachedLanes() there for consistency otherwise it's a bit odd having the cache size defined in one class and the mapping to a cache index in another.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

241

can you add support here as well? The callback is going away soon, so we also need to support the non-callback version.

llvm/lib/Transforms/Vectorize/VPlan.h

175

are you planning on adding more kinds here? otherwise this can just be a boolean? Or make this an enum class?

182

Can this have a better name, e.g. in line with the enum value or the Boolean variable name, if you change it?.

243–244

I think both should be moved to VPIteration, as we need to support both VPIteration versions there as well.

llvm/test/Transforms/LoopVectorize/AArch64/neon-extract-last-veclane.ll

3 ↗

(On Diff #319853)

This test does not need to be neon specific, right? extracting the last lane for fixed vectors should be tested fairly well already, so not sure if the test is needed at all?

llvm/test/Transforms/LoopVectorize/AArch64/sve-extract-last-veclane.ll

23

not needed?

26

nit: the names of the blocks could be improved.

[david-arm]

Hi @fhahn thanks for the review!

llvm/lib/Transforms/Vectorize/VPlan.cpp

241

OK, the reason I didn't add this originally is because I cannot test this code path in my patch. I thought it might be bad practice to add code without testing it, but I'm happy to add support here if you want. I guess we do have a test for it, so when the callback is removed if there is a bug it will break the test.

llvm/lib/Transforms/Vectorize/VPlan.h

175

I don't have plans to add other kinds here at the moment - I chose an enum here to make it extensible should people wish to add other kinds in future and based on earlier reviewer comments. I'm happy to change it to enum class or use a boolean - @sdesmalen don't know if you have a preference here?

llvm/test/Transforms/LoopVectorize/AArch64/neon-extract-last-veclane.ll

3 ↗

(On Diff #319853)

Similar to what I mentioned in the commit message, I deliberately added a llvm_unreachable() in the code that I've changed and I found no explicit tests for this at all. There is some limited coverage, but accidentally so in cases where the test was actually trying to test something else. How about I move this to a generic place?

llvm/test/Transforms/LoopVectorize/AArch64/sve-extract-last-veclane.ll

26

OK, to be honest I don't really know what they should be called. :) This is the name that LLVM generates. How about for.cond.pre-cleanup?

[fhahn]

llvm/test/Transforms/LoopVectorize/AArch64/neon-extract-last-veclane.ll

3 ↗

(On Diff #319853)

Sounds good to me

llvm/test/Transforms/LoopVectorize/AArch64/sve-extract-last-veclane.ll

26

how about something like just exit. It can also directly return; you don't need the %mul.lcssa phi I think, LV will insert them if needed.

[david-arm]

• Changed some comments in VPIteration.
• Moved getNumCachedLanes() and mapLaneToCacheIndex() to VPIteration.
• Changed enum LaneKind to enum class LaneKind, which required introducing constructors in a previous patch.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2486

I thought that adding an operator overload to struct VPIteration seemed a bit unnecessary for just this one case. The way I've defined LK_ScalableFirst to work means that the pair (Instance.Lane=0,Instance.Kind=LK_ScalableFirst) actually refers to the last element of the vector, i.e. LastLane - Instance.Lane.

llvm/lib/Transforms/Vectorize/VPlan.h

182

Hi @fhahn It's hard to come up with a much better name - how about "isKnownLane"? This terminology is used throughout the codebase to mean something that is known at compile time, e.g. ElementCount::getKnownMinValue()

241

Sadly this way I also get a compile-time warning about ending the function without returning a value, which means I still have to add a default here.

[david-arm]

• Rebase.

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2486

Okay, fair enough on not overloading the operator.

I'm a bit concerned about exposing Lane directly though, because this now means you'll need to check Instance.isKnownLane() to know whether this concerns the first lane, or the last. Any code that forgets to check this, may make the wrong assumptions.

I'd suggest making Lane private, and adding interfaces such as:

• getAsFirst(int Lane)
• getAsScalableLast(int Lane)

which asserts that the lane starts at the beginning (getFromFirst) of from the back of a scalable vector (getFromScalableLast).
It would also remove the need for additional asserts you added to check if the VPIteration Instance is a known lane.

2486

nit:

assert((!Instance.isFirstLane() ||
        Cost->isUniformAfterVectorization(cast<Instruction>(V), VF))
          && "Uniform values only have lane zero");

?

4451

Can you restructure this so that either:

• Lane/Kind are explicitly set for both Fixed and Scalable case. Or alternatively:
• Both Lane and Kind are initialized for Fixed, and only updated for scalable.

4453

What I don't like about that definition is that the representation in the LoopVectorizer is entirely different for scalable and fixed-width vectors. {0, First} and {0, LastScalable} mean the first and last element, respectively. {3, First} means the last for <4 x i32> and {3, LastScalable} means 4th last element (for <vscale x 4 x i32>). I think this is confusing and error prone. If First/ScalableLast would mean the first and last scalable "chunk" of a vector, with the index being the index as normal, then the last iteration would be described as {3, First} for <4 x i32> and {3, ScalableLast} for <vscale x 4 x i32>, and so the representation in the LoopVectorizer is more or less the same. Code-generating it is slightly different of course, but there's probably only a single place where that has to happen.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4453

OK, I'll have a think about it and see code changes would be required. I think in terms of code-generating the lane index, starting from the end is easier for the developer as the most common case will be the very last lane. That's why in my original patch I'd only catered specifically for the very last lane. If we define it in terms of chunks (really just subvectors I think) I might change some of the interfaces and names to make it clear that we're dealing with lane indices from the start of what is essentially a subvector.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2486

I'm not sure about having enum specific getAsXX functions to be honest. Having to add a new function for any new types people want to add in the future seems a bit inflexible. However, I'll try to find a way of doing something better here if possible. I do take your point about making member variables private, but that will require another NFC patch to change struct to class and adding get/set variables first. It won't be a small change. :)

[david-arm]

• Added a new VPLane class to contain all the lane offset and kind, plus add methods getKnownValue and getExpr to return the compile-time value and runtime expression, respectively.
• Changed the definition of ScalableLast to Sander's suggestion where the lane value is the offset from the start of the last subvector.
• Folded some asserts about the lane value into VPLane::getKnownValue() and VPLane::mapToCacheIndex.

[sdesmalen]

Thanks for the changes @david-arm, I think this is moving in the right direction.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2486

nit: Can you maybe write this as Instance.Lane.isFirst()? Personally I find the difference between isFirstLane and isFirstIteration a bit confusing. (same for other places)

4449–4455

is it worth creating a static VPIteration::getLastLaneForVF(ElementCount) for this?

llvm/lib/Transforms/Vectorize/VPlan.cpp

70

don't use default. Cover both cases explicitly, so that if another enum value is added, the compiler will emit a diagnostic this case is not covered.

llvm/lib/Transforms/Vectorize/VPlan.h

98

nit: if you put the comments above each enum-value, doxygen generates a nice table with a comment describing what each enum-value means.
See for example ARMLdStMultipleTiming in ARMSubtarget.h (and the generated doxygen here: https://llvm.org/doxygen/classllvm_1_1ARMSubtarget.html#ac7324b67d7e3be270177e6590f0bb1e5)

108

I'd prefer this to just be named Lane, because it is still a lane. The First or ScalableLast tells in which chunk of the vector this lane lives.

118

nit: getKnownLane.

125

nit: getLaneAsRuntimeExpr ?

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2486

Sure. I added isFirstIteration at your suggestion because it refers to Part=0 as well as Lane=0.

llvm/lib/Transforms/Vectorize/VPlan.cpp

70

OK I can do that - it just might mean adding an initialiser to Lane at the start of the function. I can't return directly from a case statement without a default as the compiler warns about functions returning void otherwise.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

70

can you just put llvm_unreachable after the switch to get rid of the warning?

[david-arm]

llvm/lib/Transforms/Vectorize/VPlan.cpp

70

Yeah I can do that if that's the preferred way?

[david-arm]

• Renamed and removed some interfaces.
• Removed default case in getAsRuntimeExpr, but had to add llvm_unreachable() as a result in order to kill the resulting compiler warning.
• Add VPLane::getLastLaneForVF().

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2942

nit: maybe add a static VPIteration::getFirstLane? (seems to be useful in several places).

llvm/lib/Transforms/Vectorize/VPlan.cpp

242

should this use mapToCacheIndex ?

llvm/lib/Transforms/Vectorize/VPlan.h

166–168

nit: const VPLane &

170–171

nit: is isFirstIteration unused now?

[david-arm]

• Rebase required changing VPTransformState since VectorizerValueMap has been removed.
• Added VPLane::getFirstLane()

[david-arm]

llvm/lib/Transforms/Vectorize/VPlan.h

170–171

It's still used in two places so I've left this in.

[sdesmalen]

Thanks for all the changes. The patch LGTM with nits addressed.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4444

nit: Did you need to move the call to Builder.SetInsertPoint?

llvm/lib/Transforms/Vectorize/VPlan.cpp

429

nit: it would be nice if you can write State->Instance->Lane = VPLane(Lane, VPLane::Kind::First);
(e.g. by providing VPLane(const VPlane &Other) in favour of having a set method).

llvm/test/Transforms/LoopVectorize/AArch64/sve-extract-last-veclane.ll

2

nit: passing the attribute in the command is redundant if it's already set by the IR attributes.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1155

It looks like this is only used in llvm/lib/Transforms/Vectorize/VPlan.cpp in this patch? Should it be defined directly there? If it needs to be shared between multiple files it would probably be better to just put the declaration into a header?

llvm/test/Transforms/LoopVectorize/extract-last-veclane.ll

56

Is there reason to use the metadata here? Can we instead just use the -force-vector-width option, which should be a bit simpler and the expected VF is clear from the run line?

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1155

It will be used by code in this file in other patches - for example, @CarolineConcatto currently has a patch that also uses this function. I can try to find a suitable header to put this in.

[fhahn]

LGTM thanks!

llvm/lib/Transforms/Vectorize/VPlan.h

63

No need for extern here I think? Please also add a comment.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 static Value *createStepForVF(IRBuilder<> &B, Constant *Step, ElementCount VF) {
   assert(isa<ConstantInt>(Step) && "Expected an integer step");
   Constant *StepVal = ConstantInt::get(
       Step->getType(),
       cast<ConstantInt>(Step)->getSExtValue() * VF.getKnownMinValue());
   return VF.isScalable() ? B.CreateVScale(StepVal) : StepVal;
 }
 
+/// Return the runtime value for VF.
+Value *getRuntimeVF(IRBuilder<> &B, Type *Ty, ElementCount VF) {
+  Constant *EC = ConstantInt::get(Ty, VF.getKnownMinValue());
+  return VF.isScalable() ? B.CreateVScale(EC) : EC;
+}
+
 namespace llvm {
 
 void reportVectorizationFailure(const StringRef DebugMsg,
     const StringRef OREMsg, const StringRef ORETag,
     OptimizationRemarkEmitter *ORE, Loop *TheLoop, Instruction *I) {
   LLVM_DEBUG(debugVectorizationFailure(DebugMsg, I));
   LoopVectorizeHints Hints(TheLoop, true /* doesn't matter */, *ORE);
   ORE->emit(createLVAnalysis(Hints.vectorizeAnalysisPassName(),

 Value *
 InnerLoopVectorizer::getOrCreateScalarValue(Value *V,
                                             const VPIteration &Instance) {
   // If the value is not an instruction contained in the loop, it should
   // already be scalar.
   if (OrigLoop->isLoopInvariant(V))
     return V;
 
-  assert(Instance.Lane > 0
+  assert(!Instance.isFirstLane()
              ? !Cost->isUniformAfterVectorization(cast<Instruction>(V), VF)
              : true && "Uniform values only have lane zero");
 
   // If the value from the original loop has not been vectorized, it is
   // represented by UF x VF scalar values in the new loop. Return the requested
   // scalar value.
   if (VectorLoopValueMap.hasScalarValue(V, Instance))
     return VectorLoopValueMap.getScalarValue(V, Instance);
 
   // If the value has not been scalarized, get its entry in VectorLoopValueMap
   // for the given unroll part. If this entry is not a vector type (i.e., the
   // vectorization factor is one), there is no need to generate an
   // extractelement instruction.
   auto *U = getOrCreateVectorValue(V, Instance.Part);
   if (!U->getType()->isVectorTy()) {
     assert(VF.isScalar() && "Value not scalarized has non-vector type");
     return U;
   }
 
+  Value *Lane = Instance.Lane.getExpr(Builder, VF);
+
   // Otherwise, the value from the original loop has been vectorized and is
   // represented by UF vector values. Extract and return the requested scalar
   // value from the appropriate vector lane.
-  return Builder.CreateExtractElement(U, Builder.getInt32(Instance.Lane));
+  return Builder.CreateExtractElement(U, Lane);
 }
 
 void InnerLoopVectorizer::packScalarIntoVectorValue(
     Value *V, const VPIteration &Instance) {
   assert(V != Induction && "The new induction variable should not be used.");
   assert(!V->getType()->isVectorTy() && "Can't pack a vector");
   assert(!V->getType()->isVoidTy() && "Type does not produce a value");
 
   Value *ScalarInst = VectorLoopValueMap.getScalarValue(V, Instance);
   Value *VectorValue = VectorLoopValueMap.getVectorValue(V, Instance.Part);
   VectorValue = Builder.CreateInsertElement(VectorValue, ScalarInst,
-                                            Builder.getInt32(Instance.Lane));
+                                            Instance.Lane.getExpr(Builder, VF));
   VectorLoopValueMap.resetVectorValue(V, Instance.Part, VectorValue);
 }
 
 void InnerLoopVectorizer::packScalarIntoVectorValue(VPValue *Def,
                                                     const VPIteration &Instance,
                                                     VPTransformState &State) {
   Value *ScalarInst = State.get(Def, Instance);
   Value *VectorValue = State.get(Def, Instance.Part);
   VectorValue = Builder.CreateInsertElement(
-      VectorValue, ScalarInst, State.Builder.getInt32(Instance.Lane));
+      VectorValue, ScalarInst, Instance.Lane.getExpr(State.Builder, VF));
   State.set(Def, VectorValue, Instance.Part);
 }
 
 Value *InnerLoopVectorizer::reverseVector(Value *Vec) {
   assert(Vec->getType()->isVectorTy() && "Invalid type");
   assert(!VF.isScalable() && "Cannot reverse scalable vectors");
   SmallVector<int, 8> ShuffleMask;
   for (unsigned i = 0; i < VF.getKnownMinValue(); ++i)
void InnerLoopVectorizer::scalarizeInstruction(Instruction *Instr, VPUser &User,
 
   // Replace the operands of the cloned instructions with their scalar
   // equivalents in the new loop.
   for (unsigned op = 0, e = User.getNumOperands(); op != e; ++op) {
     auto *Operand = dyn_cast<Instruction>(Instr->getOperand(op));
     auto InputInstance = Instance;
     if (!Operand || !OrigLoop->contains(Operand) ||
         (Cost->isUniformAfterVectorization(Operand, State.VF)))
-      InputInstance.Lane = 0;
+      InputInstance.Lane.set(0, VPLane::Kind::First);
     auto *NewOp = State.get(User.getOperand(op), InputInstance);
     Cloned->setOperand(op, NewOp);
   }
   addNewMetadata(Cloned, Instr);
 
   // Place the cloned scalar in the new loop.
   Builder.Insert(Cloned);
 
void InnerLoopVectorizer::fixLCSSAPHIs(VPTransformState &State) {
   for (PHINode &LCSSAPhi : LoopExitBlock->phis()) {
     if (LCSSAPhi.getBasicBlockIndex(LoopMiddleBlock) != -1)
       // Some phis were already hand updated by the reduction and recurrence
       // code above, leave them alone.
       continue;
 
     auto *IncomingValue = LCSSAPhi.getIncomingValue(0);
     // Non-instruction incoming values will have only one value.
-    unsigned LastLane = 0;
-    if (isa<Instruction>(IncomingValue))
-      LastLane = Cost->isUniformAfterVectorization(
-                     cast<Instruction>(IncomingValue), VF)
-                     ? 0
-                     : VF.getKnownMinValue() - 1;
-    assert((!VF.isScalable() || LastLane == 0) &&
-           "scalable vectors dont support non-uniform scalars yet");
+
+    unsigned LaneOffset = 0;
+    VPLane::Kind Kind = VPLane::Kind::First;
+    Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());
+    if (isa<Instruction>(IncomingValue) &&
+        !Cost->isUniformAfterVectorization(cast<Instruction>(IncomingValue),
+                                           VF)) {
+      LaneOffset = VF.getKnownMinValue() - 1;
+      if (VF.isScalable())
+        // In this case 'LaneOffset' refers to the offset from the start of the
+        // last subvector with VF.getKnownMinValue() elements.
+        Kind = VPLane::Kind::ScalableLast;
+      else
+        Kind = VPLane::Kind::First;
+    }
+
     // Can be a loop invariant incoming value or the last scalar value to be
     // extracted from the vectorized loop.
-    Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());
     Value *lastIncomingValue =
         OrigLoop->isLoopInvariant(IncomingValue)
             ? IncomingValue
             : State.get(State.Plan->getVPValue(IncomingValue),
-                        VPIteration(UF - 1, LastLane));
+                        VPIteration(UF - 1, LaneOffset, Kind));
     LCSSAPhi.addIncoming(lastIncomingValue, LoopMiddleBlock);
   }
 }
 
 void InnerLoopVectorizer::sinkScalarOperands(Instruction *PredInst) {
   // The basic block and loop containing the predicated instruction.
   auto *PredBB = PredInst->getParent();
   auto *VectorLoop = LI->getLoopFor(PredBB);

 void VPReplicateRecipe::execute(VPTransformState &State) {
   if (State.Instance) { // Generate a single instance.
     assert(!State.VF.isScalable() && "Can't scalarize a scalable vector");
     State.ILV->scalarizeInstruction(getUnderlyingInstr(), *this,
                                     *State.Instance, IsPredicated, State);
     // Insert scalar instance packing it into a vector.
     if (AlsoPack && State.VF.isVector()) {
       // If we're constructing lane 0, initialize to start from poison.
-      if (State.Instance->Lane == 0) {
+      if (State.Instance->isFirstLane()) {
         assert(!State.VF.isScalable() && "VF is assumed to be non scalable.");
         Value *Poison = PoisonValue::get(
             VectorType::get(getUnderlyingValue()->getType(), State.VF));
         State.ValueMap.setVectorValue(getUnderlyingInstr(),
                                       State.Instance->Part, Poison);
       }
       State.ILV->packScalarIntoVectorValue(getUnderlyingInstr(),
                                            *State.Instance);
for (unsigned Lane = 0; Lane < EndLane; ++Lane)
                                       VPIteration(Part, Lane), IsPredicated,
                                       State);
 }
 
 void VPBranchOnMaskRecipe::execute(VPTransformState &State) {
   assert(State.Instance && "Branch on Mask works only on single instance.");
 
   unsigned Part = State.Instance->Part;
-  unsigned Lane = State.Instance->Lane;
+  unsigned Lane = State.Instance->Lane.getKnownValue();
 
   Value *ConditionBit = nullptr;
   VPValue *BlockInMask = getMask();
   if (BlockInMask) {
     ConditionBit = State.get(BlockInMask, Part);
     if (ConditionBit->getType()->isVectorTy())
       ConditionBit = State.Builder.CreateExtractElement(
           ConditionBit, State.Builder.getInt32(Lane));

llvm/lib/Transforms/Vectorize/VPlan.h

 class raw_ostream;
 class RecurrenceDescriptor;
 class Value;
 class VPBasicBlock;
 class VPRegionBlock;
 class VPlan;
 class VPlanSlp;
 
 /// A range of powers-of-2 vectorization factors with fixed start and
 /// adjustable end. The range includes start and excludes end, e.g.,:
 /// [1, 9) = {1, 2, 4, 8}
 struct VFRange {
   // A power of 2.
   const ElementCount Start;
 
   // Need not be a power of 2. If End <= Start range is empty.
   ElementCount End;

 };
 
 using VPlanPtr = std::unique_ptr<VPlan>;
 
 /// In what follows, the term "input IR" refers to code that is fed into the
 /// vectorizer whereas the term "output IR" refers to code that is generated by
 /// the vectorizer.
 
+/// VPLane provides a way to access lanes in both fixed width and scalable
+/// vectors, where for the latter the lane index sometimes needs calculating
+/// as a runtime expression.
+class VPLane {
+public:
+  /// Kind describes how to interpret Val.
+  /// For First, Val is the index into the first N elements of a
+  /// fixed-vector <N x <ElTy>> or a scalable vector <vscale x N x <ElTy>>.
+  /// For ScalableLast, Val is the offset from the start of the last N-element
+  /// subvector in a scalable vector <vscale x N x <ElTy>>. For example, a Val
+  /// of 0 corresponds to lane `(vscale - 1) * N`, a Val of 1 corresponds to
+  /// `((vscale - 1) * N) + 1`, etc.
+  enum class Kind { First, ScalableLast };
+
+private:
+  /// in [0..VF)
+  unsigned Val;
+
+  /// Indicates how the Lane should be interpreted, as described above.
+  Kind LaneKind;
+
+public:
+  VPLane(unsigned Val, Kind LaneKind) : Val(Val), LaneKind(LaneKind) {}
+
+  /// Returns a compile-time known value for the lane index and asserts if the
+  /// lane can only be calculated at runtime.
+  unsigned getKnownValue() const {
+    assert(LaneKind == Kind::First);
+    return Val;
+  }
+
+  /// Returns an expression describing the lane index that can be used at
+  /// runtime.
+  Value *getExpr(IRBuilder<> &Builder, const ElementCount &VF) const;
+
+  /// Returns the Kind of lane offset.
+  Kind getKind() const { return LaneKind; }
+
+  /// Sets the lane offset and lane kind.
+  void set(unsigned V, Kind K) {
+    Val = V;
+    LaneKind = K;
+  }
+
+  /// Returns true if this is the first lane of the whole vector.
+  bool isFirst() const { return Val == 0 && LaneKind == Kind::First; }
+
+  /// Maps the lane to a cache index based on \p VF.
+  unsigned mapToCacheIndex(const ElementCount &VF) const {
+    switch (LaneKind) {
+    case VPLane::Kind::ScalableLast:
+      assert(VF.isScalable() && Val < VF.getKnownMinValue());
+      return VF.getKnownMinValue() + Val;
+    default:
+      assert(Val < VF.getKnownMinValue());
+      return Val;
+    }
+  }
+
+  /// Returns the maxmimum number of lanes that we are able to consider
+  /// caching for \p VF.
+  static unsigned getNumCachedLanes(const ElementCount &VF) {
+    return VF.getKnownMinValue() * (VF.isScalable() ? 2 : 1);
+  }
+};
+
 /// VPIteration represents a single point in the iteration space of the output
 /// (vectorized and/or unrolled) IR loop.
 struct VPIteration {
   /// in [0..UF)
   unsigned Part;
 
-  /// in [0..VF)
-  unsigned Lane;
+  VPLane Lane;
 
-  VPIteration(unsigned Part, unsigned Lane) : Part(Part), Lane(Lane) {}
+  VPIteration(unsigned Part, unsigned Lane,
+              VPLane::Kind Kind = VPLane::Kind::First)
+      : Part(Part), Lane(Lane, Kind) {}
 
-  bool isFirstIteration() const { return Part == 0 && Lane == 0; }
+  bool isFirstLane() const { return Lane.isFirst(); }
+  bool isFirstIteration() const { return Part == 0 && isFirstLane(); }
 };
 
 /// This is a helper struct for maintaining vectorization state. It's used for
 /// mapping values from the original loop to their corresponding values in
 /// the new loop. Two mappings are maintained: one for vectorized values and
 /// one for scalarized values. Vectorized values are represented with UF
 /// vector values in the new loop, and scalarized values are represented with
 /// UF x VF scalar values in the new loop. UF and VF are the unroll and
 /// vectorization factors, respectively.
 ///
 /// Entries can be added to either map with setVectorValue and setScalarValue,
 /// which assert that an entry was not already added before. If an entry is to
 /// replace an existing one, call resetVectorValue and resetScalarValue. This is
 /// currently needed to modify the mapped values during "fix-up" operations that
 /// occur once the first phase of widening is complete. These operations include
 /// type truncation and the second phase of recurrence widening.
 ///
 /// Entries from either map can be retrieved using the getVectorValue and
 /// getScalarValue functions, which assert that the desired value exists.
public:
   }
 
   /// \return True if the map has any scalar entry for \p Key.
   bool hasAnyScalarValue(Value *Key) const {
     return ScalarMapStorage.count(Key);
   }
 
   /// \return True if the map has a scalar entry for \p Key and \p Instance.
   bool hasScalarValue(Value *Key, const VPIteration &Instance) const {
     assert(Instance.Part < UF && "Queried Scalar Part is too large.");
-    assert(Instance.Lane < VF.getKnownMinValue() &&
-           "Queried Scalar Lane is too large.");
 
     if (!hasAnyScalarValue(Key))
       return false;
     const ScalarParts &Entry = ScalarMapStorage.find(Key)->second;
     assert(Entry.size() == UF && "ScalarParts has wrong dimensions.");
-    assert(Entry[Instance.Part].size() == VF.getKnownMinValue() &&
+    assert(Entry[Instance.Part].size() == VPLane::getNumCachedLanes(VF) &&
            "ScalarParts has wrong dimensions.");
-    return Entry[Instance.Part][Instance.Lane] != nullptr;
+    unsigned CacheIdx = Instance.Lane.mapToCacheIndex(VF);
+    return Entry[Instance.Part][CacheIdx] != nullptr;
   }
 
   /// Retrieve the existing vector value that corresponds to \p Key and
   /// \p Part.
   Value *getVectorValue(Value *Key, unsigned Part) {
     assert(hasVectorValue(Key, Part) && "Getting non-existent value.");
     return VectorMapStorage[Key][Part];
   }
 
   /// Retrieve the existing scalar value that corresponds to \p Key and
   /// \p Instance.
   Value *getScalarValue(Value *Key, const VPIteration &Instance) {
     assert(hasScalarValue(Key, Instance) && "Getting non-existent value.");
-    return ScalarMapStorage[Key][Instance.Part][Instance.Lane];
+    unsigned CacheIdx = Instance.Lane.mapToCacheIndex(VF);
+    return ScalarMapStorage[Key][Instance.Part][CacheIdx];
   }
 
   /// Set a vector value associated with \p Key and \p Part. Assumes such a
   /// value is not already set. If it is, use resetVectorValue() instead.
   void setVectorValue(Value *Key, unsigned Part, Value *Vector) {
     assert(!hasVectorValue(Key, Part) && "Vector value already set for part");
     if (!VectorMapStorage.count(Key)) {
       VectorParts Entry(UF);
       VectorMapStorage[Key] = Entry;
     }
     VectorMapStorage[Key][Part] = Vector;
   }
 
   /// Set a scalar value associated with \p Key and \p Instance. Assumes such a
   /// value is not already set.
   void setScalarValue(Value *Key, const VPIteration &Instance, Value *Scalar) {
     assert(!hasScalarValue(Key, Instance) && "Scalar value already set");
     if (!ScalarMapStorage.count(Key)) {
       ScalarParts Entry(UF);
       // TODO: Consider storing uniform values only per-part, as they occupy
       //       lane 0 only, keeping the other VF-1 redundant entries null.
       for (unsigned Part = 0; Part < UF; ++Part)
-        Entry[Part].resize(VF.getKnownMinValue(), nullptr);
+        Entry[Part].resize(VPLane::getNumCachedLanes(VF), nullptr);
       ScalarMapStorage[Key] = Entry;
     }
-    ScalarMapStorage[Key][Instance.Part][Instance.Lane] = Scalar;
+    unsigned CacheIdx = Instance.Lane.mapToCacheIndex(VF);
+    ScalarMapStorage[Key][Instance.Part][CacheIdx] = Scalar;
   }
 
   /// Reset the vector value associated with \p Key for the given \p Part.
   /// This function can be used to update values that have already been
   /// vectorized. This is the case for "fix-up" operations including type
   /// truncation and the second phase of recurrence vectorization.
   void resetVectorValue(Value *Key, unsigned Part, Value *Vector) {
     assert(hasVectorValue(Key, Part) && "Vector value not set for part");
     VectorMapStorage[Key][Part] = Vector;
   }
 
   /// Reset the scalar value associated with \p Key for \p Part and \p Lane.
   /// This function can be used to update values that have already been
   /// scalarized. This is the case for "fix-up" operations including scalar phi
   /// nodes for scalarized and predicated instructions.
   void resetScalarValue(Value *Key, const VPIteration &Instance,
                         Value *Scalar) {
     assert(hasScalarValue(Key, Instance) &&
            "Scalar value not set for part and lane");
-    ScalarMapStorage[Key][Instance.Part][Instance.Lane] = Scalar;
+    unsigned CacheIdx = Instance.Lane.mapToCacheIndex(VF);
+    ScalarMapStorage[Key][Instance.Part][CacheIdx] = Scalar;
   }
 };
 
 /// This class is used to enable the VPlan to invoke a method of ILV. This is
 /// needed until the method is refactored out of ILV and becomes reusable.
 struct VPCallback {
   virtual ~VPCallback() {}
   virtual Value *getOrCreateVectorValues(Value *V, unsigned Part) = 0;
return I != Data.PerPartOutput.end() && Part < I->second.size() &&
            I->second[Part];
   }
 
   bool hasScalarValue(VPValue *Def, VPIteration Instance) {
     auto I = Data.PerPartScalars.find(Def);
     if (I == Data.PerPartScalars.end())
       return false;
     return Instance.Part < I->second.size() &&
-           Instance.Lane < I->second[Instance.Part].size() &&
-           I->second[Instance.Part][Instance.Lane];
+           Instance.Lane.getKnownValue() < I->second[Instance.Part].size() &&
+           I->second[Instance.Part][Instance.Lane.getKnownValue()];
   }
 
   /// Set the generated Value for a given VPValue and a given Part.
   void set(VPValue *Def, Value *V, unsigned Part) {
     if (!Data.PerPartOutput.count(Def)) {
       DataState::PerPartValuesTy Entry(UF);
       Data.PerPartOutput[Def] = Entry;
     }
     Data.PerPartOutput[Def][Part] = V;
   }
   void set(VPValue *Def, Value *IRDef, Value *V, unsigned Part);
   void reset(VPValue *Def, Value *IRDef, Value *V, unsigned Part);
   void set(VPValue *Def, Value *IRDef, Value *V, const VPIteration &Instance);
 
   void set(VPValue *Def, Value *V, const VPIteration &Instance) {
     auto Iter = Data.PerPartScalars.insert({Def, {}});
     auto &PerPartVec = Iter.first->second;
     while (PerPartVec.size() <= Instance.Part)
       PerPartVec.emplace_back();
     auto &Scalars = PerPartVec[Instance.Part];
-    while (Scalars.size() <= Instance.Lane)
+    while (Scalars.size() <= Instance.Lane.getKnownValue())
       Scalars.push_back(nullptr);
-    Scalars[Instance.Lane] = V;
+    Scalars[Instance.Lane.getKnownValue()] = V;
   }
 
   /// Hold state information used when constructing the CFG of the output IR,
   /// traversing the VPBasicBlocks and generating corresponding IR BasicBlocks.
   struct CFGState {
     /// The previous VPBasicBlock visited. Initially set to null.
     VPBasicBlock *PrevVPBB = nullptr;
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <cassert>
 #include <iterator>
 #include <string>
 #include <vector>
 
 using namespace llvm;
 extern cl::opt<bool> EnableVPlanNativePath;
+extern Value *getRuntimeVF(IRBuilder<> &B, Type *Ty, ElementCount VF);
 
 #define DEBUG_TYPE "vplan"
 
 raw_ostream &llvm::operator<<(raw_ostream &OS, const VPValue &V) {
   const VPInstruction *Instr = dyn_cast<VPInstruction>(&V);
   VPSlotTracker SlotTracker(
       (Instr && Instr->getParent()) ? Instr->getParent()->getPlan() : nullptr);
   V.print(OS, SlotTracker);
   return OS;
 }
 
+Value *VPLane::getExpr(IRBuilder<> &Builder, const ElementCount &VF) const {
+  Value *Lane;
+  switch (LaneKind) {
+  case VPLane::Kind::ScalableLast:
+    // Lane = RuntimeVF - VF.getKnownMinValue() + Val
+    Lane = Builder.CreateSub(getRuntimeVF(Builder, Builder.getInt32Ty(), VF),
+                             Builder.getInt32(VF.getKnownMinValue() - Val));
+    break;
+  default:
+    assert(LaneKind == VPLane::Kind::First);
+    Lane = Builder.getInt32(Val);
+    break;
+  }
+  return Lane;
+}
+
 VPValue::VPValue(const unsigned char SC, Value *UV, VPDef *Def)
     : SubclassID(SC), UnderlyingVal(UV), Def(Def) {
   if (Def)
     Def->addDefinedValue(this);
 }
 
 VPValue::~VPValue() {
   assert(Users.empty() && "trying to delete a VPValue with remaining users");
VPBasicBlock::iterator VPBasicBlock::getFirstNonPhi() {
   return It;
 }
 
 Value *VPTransformState::get(VPValue *Def, const VPIteration &Instance) {
   if (!Def->getDef())
     return Def->getLiveInIRValue();
 
   if (hasScalarValue(Def, Instance))
-    return Data.PerPartScalars[Def][Instance.Part][Instance.Lane];
+    return Data
+        .PerPartScalars[Def][Instance.Part][Instance.Lane.getKnownValue()];
 
   if (hasVectorValue(Def, Instance.Part)) {
     assert(Data.PerPartOutput.count(Def));
     auto *VecPart = Data.PerPartOutput[Def][Instance.Part];
     if (!VecPart->getType()->isVectorTy()) {
-      assert(Instance.Lane == 0 && "cannot get lane > 0 for scalar");
+      assert(Instance.isFirstLane() && "cannot get lane > 0 for scalar");
       return VecPart;
     }
+
+    Value *Lane = Instance.Lane.getExpr(Builder, VF);
     // TODO: Cache created scalar values.
-    return Builder.CreateExtractElement(VecPart,
-                                        Builder.getInt32(Instance.Lane));
+    return Builder.CreateExtractElement(VecPart, Lane);
   }
   return Callback.getOrCreateScalarValue(VPValue2Value[Def], Instance);
 }
 
 BasicBlock *
 VPBasicBlock::createEmptyBasicBlock(VPTransformState::CFGState &CFG) {
   // BB stands for IR BasicBlocks. VPBB stands for VPlan VPBasicBlocks.
   // Pred stands for Predessor. Prev stands for Previous - last visited/created.
void VPRegionBlock::execute(VPTransformState *State) {
   // Enter replicating mode.
   State->Instance = VPIteration(0, 0);
 
   for (unsigned Part = 0, UF = State->UF; Part < UF; ++Part) {
     State->Instance->Part = Part;
     assert(!State->VF.isScalable() && "VF is assumed to be non scalable.");
     for (unsigned Lane = 0, VF = State->VF.getKnownMinValue(); Lane < VF;
          ++Lane) {
-      State->Instance->Lane = Lane;
+      State->Instance->Lane.set(Lane, VPLane::Kind::First);
       // Visit the VPBlocks connected to \p this, starting from it.
       for (VPBlockBase *Block : RPOT) {
         LLVM_DEBUG(dbgs() << "LV: VPBlock in RPO " << Block->getName() << '\n');
         Block->execute(State);
       }
     }
   }
 

llvm/test/Transforms/LoopVectorize/AArch64/sve-extract-last-veclane.ll

+; RUN: opt -loop-vectorize -dce -instcombine -mtriple aarch64-linux-gnu -mattr=+sve -S < %s 2>%t | FileCheck %s
+
+; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
+
+; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
+; WARN-NOT: warning
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define void @inv_store_last_lane(i32* noalias nocapture %a, i32* noalias nocapture %inv, i32* noalias nocapture readonly %b, i64 %n) #0 {
+; CHECK-LABEL: @inv_store_last_lane
+; CHECK: vector.body:
+; CHECK:  store <vscale x 4 x i32> %[[VEC_VAL:.*]], <
+; CHECK: middle.block:
+; CHECK: %[[VSCALE:.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT: %[[VSCALE2:.*]] = shl i32 %[[VSCALE]], 2
+; CHECK-NEXT: %[[LAST_LANE:.*]] = add i32 %[[VSCALE2]], -1
+; CHECK-NEXT: %{{.*}} = extractelement <vscale x 4 x i32> %[[VEC_VAL]], i32 %[[LAST_LANE]]
+
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.body
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %mul = shl nsw i32 %0, 1
+  %arrayidx2 = getelementptr inbounds i32, i32* %a, i64 %indvars.iv
+  store i32 %mul, i32* %arrayidx2, 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 %exit, label %for.body, !llvm.loop !0
+
+exit:              ; preds = %for.body
+  %arrayidx5 = getelementptr inbounds i32, i32* %inv, i64 42
+  store i32 %mul, i32* %arrayidx5, align 4
+  ret void
+}
+
+define float @ret_last_lane(float* noalias nocapture %a, float* noalias nocapture readonly %b, i64 %n) #0 {
+; CHECK-LABEL: @ret_last_lane
+; CHECK: vector.body:
+; CHECK:  store <vscale x 4 x float> %[[VEC_VAL:.*]], <
+; CHECK: middle.block:
+; CHECK: %[[VSCALE:.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT: %[[VSCALE2:.*]] = shl i32 %[[VSCALE]], 2
+; CHECK-NEXT: %[[LAST_LANE:.*]] = add i32 %[[VSCALE2]], -1
+; CHECK-NEXT: %{{.*}} = extractelement <vscale x 4 x float> %[[VEC_VAL]], i32 %[[LAST_LANE]]
+
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %b, i64 %indvars.iv
+  %0 = load float, float* %arrayidx, align 4
+  %mul = fmul float %0, 2.000000e+00
+  %arrayidx2 = getelementptr inbounds float, float* %a, i64 %indvars.iv
+  store float %mul, float* %arrayidx2, 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 %exit, label %for.body, !llvm.loop !6
+
+exit:                                 ; preds = %for.body, %entry
+  ret float %mul
+}
+
+attributes #0 = { "target-cpu"="generic" "target-features"="+neon,+sve" }
+
+!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}
+!6 = distinct !{!6, !1, !2, !3, !4, !5}

llvm/test/Transforms/LoopVectorize/extract-last-veclane.ll

+; RUN: opt -loop-vectorize -dce -instcombine -S < %s 2>%t | FileCheck %s
+
+define void @inv_store_last_lane(i32* noalias nocapture %a, i32* noalias nocapture %inv, i32* noalias nocapture readonly %b, i64 %n) {
+; CHECK-LABEL: @inv_store_last_lane
+; CHECK: vector.body:
+; CHECK:  store <4 x i32> %[[VEC_VAL:.*]], <
+; CHECK: middle.block:
+; CHECK:  %{{.*}} = extractelement <4 x i32> %[[VEC_VAL]], i32 3
+
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %mul = shl nsw i32 %0, 1
+  %arrayidx2 = getelementptr inbounds i32, i32* %a, i64 %indvars.iv
+  store i32 %mul, i32* %arrayidx2, 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 %exit, label %for.body, !llvm.loop !0
+
+exit:              ; preds = %for.body
+  %arrayidx5 = getelementptr inbounds i32, i32* %inv, i64 42
+  store i32 %mul, i32* %arrayidx5, align 4
+  ret void
+}
+
+define float @ret_last_lane(float* noalias nocapture %a, float* noalias nocapture readonly %b, i64 %n) {
+; CHECK-LABEL: @ret_last_lane
+; CHECK: vector.body:
+; CHECK:  store <4 x float> %[[VEC_VAL:.*]], <
+; CHECK: middle.block:
+; CHECK:  %{{.*}} = extractelement <4 x float> %[[VEC_VAL]], i32 3
+
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %b, i64 %indvars.iv
+  %0 = load float, float* %arrayidx, align 4
+  %mul = fmul float %0, 2.000000e+00
+  %arrayidx2 = getelementptr inbounds float, float* %a, i64 %indvars.iv
+  store float %mul, float* %arrayidx2, 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 %exit, label %for.body, !llvm.loop !6
+
+exit:                                 ; preds = %for.body, %entry
+  ret float %mul
+}
+
+!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 false}
+!4 = !{!"llvm.loop.interleave.count", i32 1}
+!5 = !{!"llvm.loop.vectorize.enable", i1 true}
+!6 = distinct !{!6, !1, !2, !3, !4, !5}