[llvm] [VPlan] Unroll VPReplicateRecipe by VF. (PR #142433)

[Florian Hahn via llvm-commits]

================
@@ -2626,43 +2664,29 @@ static void scalarizeInstruction(const Instruction *Instr,
 
 void VPReplicateRecipe::execute(VPTransformState &State) {
   Instruction *UI = getUnderlyingInstr();
-  if (State.Lane) { // Generate a single instance.
-    assert((State.VF.isScalar() || !isSingleScalar()) &&
-           "uniform recipe shouldn't be predicated");
-    assert(!State.VF.isScalable() && "Can't scalarize a scalable vector");
-    scalarizeInstruction(UI, this, *State.Lane, State);
-    // Insert scalar instance packing it into a vector.
-    if (State.VF.isVector() && shouldPack()) {
-      // If we're constructing lane 0, initialize to start from poison.
-      if (State.Lane->isFirstLane()) {
-        assert(!State.VF.isScalable() && "VF is assumed to be non scalable.");
-        Value *Poison =
-            PoisonValue::get(VectorType::get(UI->getType(), State.VF));
-        State.set(this, Poison);
-      }
-      State.packScalarIntoVectorizedValue(this, *State.Lane);
-    }
-    return;
-  }
 
-  if (IsSingleScalar) {
-    // Uniform within VL means we need to generate lane 0.
+  if (!State.Lane) {
+    assert(IsSingleScalar &&
+           "VPReplicateRecipes outside replicate regions must be unrolled");
     scalarizeInstruction(UI, this, VPLane(0), State);
     return;
   }
 
-  // A store of a loop varying value to a uniform address only needs the last
-  // copy of the store.
-  if (isa<StoreInst>(UI) && vputils::isSingleScalar(getOperand(1))) {
-    auto Lane = VPLane::getLastLaneForVF(State.VF);
-    scalarizeInstruction(UI, this, VPLane(Lane), State);
-    return;
+  assert((State.VF.isScalar() || !isSingleScalar()) &&
+         "uniform recipe shouldn't be predicated");
+  assert(!State.VF.isScalable() && "Can't scalarize a scalable vector");
+  scalarizeInstruction(UI, this, *State.Lane, State);
+  // Insert scalar instance packing it into a vector.
+  if (State.VF.isVector() && shouldPack()) {
+    // If we're constructing lane 0, initialize to start from poison.
+    if (State.Lane->isFirstLane()) {
+      assert(!State.VF.isScalable() && "VF is assumed to be non scalable.");
+      Value *Poison =
+          PoisonValue::get(VectorType::get(UI->getType(), State.VF));
+      State.set(this, Poison);
+    }
+    State.packScalarIntoVectorizedValue(this, *State.Lane);
   }
----------------
fhahn wrote:

Sounds like this should work as well.

I sketched it here: https://github.com/llvm/llvm-project/pull/145188. When unrolling by VF, we keep a mapping of defs -> lane definintions and use that to lower/serve Pack/Unpack VPInstructions. 

It still needs some cleanup and we still have to generate some extracts.

Here's the code for Pack:https://github.com/llvm/llvm-project/pull/145188/commits/51b88fa31bebe2f6efb2e36897daaa30d785f00d 

And unpack: https://github.com/llvm/llvm-project/pull/145188/commits/af03b31f1e5c21c3382d3264e984cb10a5f749b5

And we need also need to pack replicate recipes in a vector, if they are used by replicate regions, so the VPReplicateRecipes in replicate regions can access the unrolled values: https://github.com/llvm/llvm-project/pull/145188/commits/7f53c1cea717f7479330ff7f10f7648f585b8555

Given that this requires quite a few improvements, might be simpler to landing initial unrolling first, then the pack/unpack patches?

https://github.com/llvm/llvm-project/pull/142433