[VPlan] VPTransformState::get() can always return lane 0 for uniforms.

[fhahn]

Summary

When requesting a scalar value for a uniform VPDef, we can always return
lane 0. This can avoid unnecessary inserting some unncessary instructions
to duplicate the uniform value across lanes.

[reames]

LGTM - though, please keep in mind I'm not fully familiar with this code. You may want to wait for another reviewer.

[reames]

Actually, LGTM revoked. This apparently depends on D91500 (which isn't marked in the metadata), and as commented there I don't have context.

[a.elovikov]

Not blocking this review, but I think it's bug-prone to mix lane 0 of scalarized divergent values and truly uniform values that can be kept on a single scalar. Possible examples:

top-loop:
  if %iv % VF != 0:
    inner-loop:
      %iv = [ 0, inner.ph ], [ %iv.next, inner.latch ] ; Uniform, but lane0 doesn't make much sense since it masked out.
      ...
      divergent exit condition

bb:
  %sel = select i1 %divergent, 42, %divergent.def ; divergent in general
  use %sel
  br i1 %divergent, label %uni.use.bb, label %bb2

uni.use.bb:
  %uni.phi = phi [ %sel, %bb ] ; "Conditionally" uniform - all active lanes have the same uniform value
  ; Long compute chain based on %uni.phi that we'd like to keep on a single scalar

In the latter case the correct extract for the uniform value would be from the first *active* lane, not from the lane 0. And I believe it's very easy to make a mistake if the same data storage is used for both scalarized parts of divergent values and for really uniform values that should be kept on a single scalar def/register.

To summarize - I think it's possible to implement everything correctly by repurposing lane0 storage for keeping uniform values, but it might lead (in future, once we try to implement more complex/complicated optimizations) to unexpected confusions and omissions that might lead to silent miscompiles (e.g. extracting undef values from lane0 instead of extracting required uniform values from the first active lane).

llvm/lib/Transforms/Vectorize/VPlan.h

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.
-  Value *get(VPValue *Def, const VPIteration &Instance) {
-    // If the Def is managed directly by VPTransformState, extract the lane from
-    // the relevant part. Note that currently only VPInstructions and external
-    // defs are managed by VPTransformState. Other Defs are still created by ILV
-    // and managed in its ValueMap. For those this method currently just
-    // delegates the call to ILV below.
-    if (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");
-        return VecPart;
-      }
-      // TODO: Cache created scalar values.
-      return Builder.CreateExtractElement(VecPart,
-                                          Builder.getInt32(Instance.Lane));
-    }
-
-    return Callback.getOrCreateScalarValue(VPValue2Value[Def], Instance);
-  }
+  Value *get(VPValue *Def, VPIteration 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;
     }
     Data.PerPartOutput[Def][Part] = V;

llvm/lib/Transforms/Vectorize/VPlan.cpp

void VPBlockBase::deleteCFG(VPBlockBase *Entry) {
 
   for (VPBlockBase *Block : Blocks)
     delete Block;
 }
 
 VPBasicBlock::iterator VPBasicBlock::getFirstNonPhi() {
   iterator It = begin();
   while (It != end() && (isa<VPWidenPHIRecipe>(&*It) ||
                          isa<VPWidenIntOrFpInductionRecipe>(&*It) ||
                          isa<VPPredInstPHIRecipe>(&*It) ||
                          isa<VPWidenCanonicalIVRecipe>(&*It)))
     It++;
   return It;
 }
 
+Value *VPTransformState::get(VPValue *Def, VPIteration Instance) {
+  // For uniform definitions, all lanes produce the same value, so we can always
+  // return the first lane.
+  if (auto *ReplicateR = dyn_cast<VPReplicateRecipe>(Def))
+    if (ReplicateR->isUniform())
+      Instance.Lane = 0;
+  // If the Def is managed directly by VPTransformState, extract the lane from
+  // the relevant part. Note that currently only VPInstructions and external
+  // defs are managed by VPTransformState. Other Defs are still created by ILV
+  // and managed in its ValueMap. For those this method currently just
+  // delegates the call to ILV below.
+  if (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");
+      return VecPart;
+    }
+    // TODO: Cache created scalar values.
+    return Builder.CreateExtractElement(VecPart,
+                                        Builder.getInt32(Instance.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.
   BasicBlock *PrevBB = CFG.PrevBB;
   BasicBlock *NewBB = BasicBlock::Create(PrevBB->getContext(), getName(),
                                          PrevBB->getParent(), CFG.LastBB);
   LLVM_DEBUG(dbgs() << "LV: created " << NewBB->getName() << '\n');