[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

102

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

102

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

2469

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

2490–2496

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.

4432

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

llvm/lib/Transforms/Vectorize/VPlan.h

101–108

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>>

102

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>

201

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

208

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

212–213

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

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4432

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

101–108

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.

212–213

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

225

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

110

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

117

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

212–213

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

225

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

110

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

2469

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

117

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()

208

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

2469

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.

2469

nit:

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

?

4430

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.

4432

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

4432

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

2469

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

2469

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)

4428–4434

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

2469

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

2925–2927

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

llvm/lib/Transforms/Vectorize/VPlan.cpp

226

should this use mapToCacheIndex ?

llvm/lib/Transforms/Vectorize/VPlan.h

101–107

nit: const VPLane &

109–138

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

109–138

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

4423

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

llvm/lib/Transforms/Vectorize/VPlan.cpp

423

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

1147

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

1147

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.Lane > 0 || !Instance.isKnownLane())
              ? !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;
+  switch (Instance.Kind) {
+  case VPIteration::LaneKind::LK_First:
+    Lane = Builder.getInt32(Instance.Lane);
+    break;
+  case VPIteration::LaneKind::LK_ScalableLast:
+    Lane = Builder.CreateSub(getRuntimeVF(Builder, Builder.getInt32Ty(), VF),
+                             Builder.getInt32(1 + Instance.Lane));
+    break;
+  }
+
   // 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);
+  assert(Instance.isKnownLane() && "Only constant lane indices supported");
   VectorValue = Builder.CreateInsertElement(VectorValue, ScalarInst,
                                             Builder.getInt32(Instance.Lane));
   VectorLoopValueMap.resetVectorValue(V, Instance.Part, VectorValue);
 }
 
 Value *InnerLoopVectorizer::reverseVector(Value *Vec) {
   assert(Vec->getType()->isVectorTy() && "Invalid type");
   assert(!VF.isScalable() && "Cannot reverse scalable vectors");
if (!IsVoidRetTy)
     Cloned->setName(Instr->getName() + ".cloned");
 
   // 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)))
+        (Cost->isUniformAfterVectorization(Operand, State.VF))) {
       InputInstance.Lane = 0;
+      InputInstance.Kind = VPIteration::LaneKind::LK_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() {
   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 Lane = 0;
+    VPIteration::LaneKind Kind = VPIteration::LaneKind::LK_First;
+    Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());
+    if (isa<Instruction>(IncomingValue) &&
+        !Cost->isUniformAfterVectorization(cast<Instruction>(IncomingValue),
+                                           VF)) {
+      if (VF.isScalable())
+        // In this case 'Lane' refers to the lane offset from the last lane
+        Kind = VPIteration::LaneKind::LK_ScalableLast;
+      else
+        Lane = VF.getKnownMinValue() - 1;
+    }
+
     // 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 =
-        getOrCreateScalarValue(IncomingValue, VPIteration(UF - 1, LastLane));
+        getOrCreateScalarValue(IncomingValue, VPIteration(UF - 1, Lane, 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);

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;

 /// the vectorizer.
 
 /// 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;
 
-  VPIteration() : Part(0), Lane(0) {}
-  VPIteration(unsigned Part, unsigned Lane) : Part(Part), Lane(Lane) {}
-
-  bool isNull() const { return Part == 0 && Lane == 0; }
+  /// LaneKind describes how to interpret Lane.
+  /// For LK_First, Lane 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 LK_ScalableLast, Lane is the offset from the last lane of a scalable
+  /// vector <vscale x N x  <ElTy>>. For example, a Lane of 0 corresponds to
+  /// the last lane, a Lane of 1 means last lane - 1, etc.
+  enum class LaneKind { LK_First, LK_ScalableLast };
+
+  /// Indicates how the Lane should be interpreted, as described above.
+  LaneKind Kind;
+
+  VPIteration() : Part(0), Lane(0), Kind(LaneKind::LK_First) {}
+  VPIteration(unsigned Part, unsigned Lane, LaneKind Kind = LaneKind::LK_First)
+      : Part(Part), Lane(Lane), Kind(Kind) {}
+
+  bool isNull() const {
+    return Part == 0 && Lane == 0 && Kind == LaneKind::LK_First;
+  }
+
+  bool isKnownLane() const { return Kind == LaneKind::LK_First; }
+
+  /// 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);
+  }
+
+  /// Maps the Lane to a cache index based on \p VF.
+  unsigned mapLaneToIndex(const ElementCount &VF) const {
+    switch (Kind) {
+    case LaneKind::LK_ScalableLast:
+      assert(VF.isScalable());
+      return VF.getKnownMinValue() + Lane;
+    default:
+      assert(Kind == LaneKind::LK_First);
+      return Lane;
+    }
+  }
 };
 
 /// 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
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() == VPIteration::getNumCachedLanes(VF) &&
            "ScalarParts has wrong dimensions.");
-    return Entry[Instance.Part][Instance.Lane] != nullptr;
+    unsigned CacheIdx = Instance.mapLaneToIndex(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.mapLaneToIndex(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(VPIteration::getNumCachedLanes(VF), nullptr);
       ScalarMapStorage[Key] = Entry;
     }
-    ScalarMapStorage[Key][Instance.Part][Instance.Lane] = Scalar;
+    unsigned CacheIdx = Instance.mapLaneToIndex(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.mapLaneToIndex(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;
bool hasVectorValue(VPValue *Def, unsigned Part) {
     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;
+    assert(Instance.isKnownLane() && "Non-constant lanes unsupported");
     return Instance.Part < I->second.size() &&
            Instance.Lane < I->second[Instance.Part].size() &&
            I->second[Instance.Part][Instance.Lane];
   }
 
   /// 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 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];
+    assert(Instance.isKnownLane() && "Non-constant lanes unsupported");
     while (Scalars.size() <= Instance.Lane)
       Scalars.push_back(nullptr);
     Scalars[Instance.Lane] = V;
   }
 
   /// Hold state information used when constructing the CFG of the output IR,
   /// traversing the VPBasicBlocks and generating corresponding IR BasicBlocks.
   struct CFGState {

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;
 }
 
 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");
   if (Def)
     Def->removeDefinedValue(this);
 }
 
 void VPValue::print(raw_ostream &OS, VPSlotTracker &SlotTracker) const {
   if (const VPRecipeBase *R = dyn_cast_or_null<VPRecipeBase>(Def))
     R->print(OS, "", SlotTracker);
   else
     printAsOperand(OS, SlotTracker);
VPBasicBlock::iterator VPBasicBlock::getFirstNonPhi() {
   return It;
 }
 
 Value *VPTransformState::get(VPValue *Def, const VPIteration &Instance) {
   if (!Def->getDef() && OrigLoop->isLoopInvariant(Def->getLiveInIRValue()))
     return Def->getLiveInIRValue();
 
   if (hasScalarValue(Def, Instance))
     return Data.PerPartScalars[Def][Instance.Part][Instance.Lane];
 
   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.Lane == 0 &&
+             Instance.Kind == VPIteration::LaneKind::LK_First &&
+             "cannot get lane > 0 for scalar");
       return VecPart;
     }
+
+    Value *Lane;
+    switch (Instance.Kind) {
+    case VPIteration::LaneKind::LK_First:
+      Lane = Builder.getInt32(Instance.Lane);
+      break;
+    case VPIteration::LaneKind::LK_ScalableLast:
+      Lane = Builder.CreateSub(getRuntimeVF(Builder, Builder.getInt32Ty(), VF),
+                               Builder.getInt32(1 + Instance.Lane));
+      break;
+    }
     // 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();
 
   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;
       // 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}