[VPlan] Support extracting lanes for defs managed in VPTransformState.

[fhahn]

Summary

Currently extracting a lane for a VPValue def is not supported, if it is
managed directly by VPTransformState (e.g. because it is created by a
VPInstruction or an external VPValue def).

For now, simply extract the requested lane. In the future, we should
also cache the extracted scalar values, similar to LV.

[SjoerdMeijer]

llvm/lib/Transforms/Vectorize/VPlan.h

283

Probably the comment about falling back to ILV now needs to be placed here?

And perhaps a brief comment added to the lines/comments above, what the reason could be for a Value not having a Def (at least I would be curious about that)?

[fhahn]

Adjust comments, clarify which Defs are directly managed by VPTransformState and that we delegate to ILV for others.

[SjoerdMeijer]

Thanks for this. I have played quite a bit with this patch (in the context of D79100), and this behaves entirely as expected. The change is minimal/straightforward, so LGTM.

llvm/lib/Transforms/Vectorize/VPlan.h

278

Really minor nit: perhaps "For those" -> "For those Defs"
(think that reads a bit easier)

[Ayal]

llvm/lib/Transforms/Vectorize/VPlan.h

281

This should in general also support scalar Values set per part, i.e., invariants.
Going forward, better model extracts and inserts upfront in VPlan, rather than generating them on-demand during code-gen.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.h

281

Going forward, better model extracts and inserts upfront in VPlan, rather than generating them on-demand during code-gen.

Sounds good! I'll land the patch as-is tomorrow as a first step, unless there's additional feedback.

llvm/lib/Transforms/Vectorize/VPlan.h

struct VPTransformState {
   Value *get(VPValue *Def, unsigned Part) {
     // If Values have been set for this Def return the one relevant for \p Part.
     if (Data.PerPartOutput.count(Def))
       return Data.PerPartOutput[Def][Part];
     // Def is managed by ILV: bring the Values from ValueMap.
     return Callback.getOrCreateVectorValues(VPValue2Value[Def], Part);
   }
 
-  /// Get the generated Value for a given VPValue and given Part and Lane. Note
-  /// that as per-lane Defs are still created by ILV and managed in its ValueMap
-  /// this method currently just delegates the call to ILV.
+  /// Get the generated Value for a given VPValue and given Part and Lane.
   Value *get(VPValue *Def, const VPIteration &Instance) {
+    // If values have been set for this Def, extract the lane from the relevant
+    // part.
+    // TODO: Cache created scalar values.
+    if (Data.PerPartOutput.count(Def)) {
+      auto *VecPart = Data.PerPartOutput[Def][Instance.Part];
+      return Builder.CreateExtractElement(VecPart,
+                                          Builder.getInt32(Instance.Lane));
+    }
+
     return Callback.getOrCreateScalarValue(VPValue2Value[Def], Instance);
   }
 
   /// 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;