[llvm] [VPlan] Implement interleaving as VPlan-to-VPlan transform. (PR #95842)

[Florian Hahn via llvm-commits]

================
@@ -1622,3 +1622,398 @@ void VPlanTransforms::dropPoisonGeneratingRecipes(
     }
   }
 }
+
+namespace {
+class InterleaveState {
+  DenseMap<VPValue *, SmallVector<VPValue *>> InterleavedValues;
+
+public:
+  VPValue *getInterleavedValue(VPValue *V, unsigned IC) {
+    if (IC == 0 || V->isLiveIn())
+      return V;
+    assert(
+        (InterleavedValues.contains(V) && InterleavedValues[V].size() >= IC) &&
+        "accessed value does not exist");
+    return InterleavedValues[V][IC - 1];
+  }
+
+  void addInterleavedValues(VPRecipeBase *OrigR, VPRecipeBase *CopyR,
+                            unsigned I) {
+    for (const auto &[Idx, VPV] : enumerate(OrigR->definedValues())) {
+      auto Ins = InterleavedValues.insert({VPV, {}});
+      assert(Ins.first->second.size() == I - 1 && "earlier parts not set");
+      Ins.first->second.push_back(CopyR->getVPValue(Idx));
+    }
+  }
+
+  void addUniform(VPSingleDefRecipe *R, unsigned IC) {
+    auto Ins = InterleavedValues.insert({R, {}});
+    for (unsigned I = 1; I != IC; ++I)
+      Ins.first->second.push_back(R);
+  }
+
+  bool contains(VPValue *VPV) { return InterleavedValues.contains(VPV); }
+
+  DenseMap<VPValue *, SmallVector<VPValue *>> &getInterleavedValues() {
+    return InterleavedValues;
+  }
+
+  void remapOperand(VPRecipeBase *R, unsigned OpIdx, unsigned Part) {
+    auto *Op = R->getOperand(OpIdx);
+    R->setOperand(OpIdx, getInterleavedValue(Op, Part));
+  }
+
+  void remapOperands(VPRecipeBase *R, unsigned Part) {
+    for (const auto &[OpIdx, Op] : enumerate(R->operands()))
+      R->setOperand(OpIdx, getInterleavedValue(Op, Part));
+  }
+};
+} // namespace
+
+static void interleaveReplicateRegion(VPRegionBlock *VPR, VPlan &Plan,
+                                      unsigned IC,
+                                      InterleaveState &InterleavedValues) {
+  Type *CanIVIntTy = Plan.getCanonicalIV()->getScalarType();
+  VPBlockBase *InsertPt = VPR->getSingleSuccessor();
+  for (unsigned I = 1; I != IC; ++I) {
+    auto *Copy = VPR->clone();
+    VPBlockUtils::insertBlockBefore(Copy, InsertPt);
+
+    auto PartI = vp_depth_first_shallow(Copy->getEntry());
+    auto Part0 = vp_depth_first_shallow(VPR->getEntry());
+    for (const auto &[PartIVPBB, Part0VPBB] :
+         zip(VPBlockUtils::blocksOnly<VPBasicBlock>(PartI),
+             VPBlockUtils::blocksOnly<VPBasicBlock>(Part0))) {
+      for (const auto &[PartIR, Part0R] : zip(*PartIVPBB, *Part0VPBB)) {
+        InterleavedValues.remapOperands(&PartIR, I);
+        if (auto *ScalarIVSteps = dyn_cast<VPScalarIVStepsRecipe>(&PartIR)) {
+          ScalarIVSteps->addOperand(
+              Plan.getOrAddLiveIn(ConstantInt::get(CanIVIntTy, I)));
+        }
+
+        InterleavedValues.addInterleavedValues(&Part0R, &PartIR, I);
+      }
+    }
+  }
+}
+
+static void interleaveWidenInduction(VPWidenIntOrFpInductionRecipe *IV,
+                                     VPlan &Plan, unsigned IC,
+                                     VPBasicBlock::iterator &InsertPtForPhi,
+                                     InterleaveState &InterleavedValues,
+                                     VPTypeAnalysis &TypeInfo,
+                                     SmallPtrSet<VPRecipeBase *, 8> &ToSkip) {
+  VPBasicBlock *PH = cast<VPBasicBlock>(
+      IV->getParent()->getEnclosingLoopRegion()->getSinglePredecessor());
+  VPValue *Step = &Plan.getVF();
+  Type *IVTy = TypeInfo.inferScalarType(IV);
+  auto &ID = IV->getInductionDescriptor();
+  FastMathFlags FMFs;
+  if (ID.getInductionBinOp() && isa<FPMathOperator>(ID.getInductionBinOp()))
+    FMFs = ID.getInductionBinOp()->getFastMathFlags();
+
+  if (TypeInfo.inferScalarType(Step) != IVTy) {
+    Instruction::CastOps CastOp =
+        IVTy->isFloatingPointTy() ? Instruction::UIToFP : Instruction::Trunc;
+    Step = new VPWidenCastRecipe(CastOp, Step, IV->getScalarType());
+    PH->appendRecipe(Step->getDefiningRecipe());
+    ToSkip.insert(Step->getDefiningRecipe());
+  }
+
+  auto *ConstScale =
+      IV->getOperand(1)->isLiveIn()
+          ? dyn_cast<ConstantInt>(IV->getOperand(1)->getLiveInIRValue())
+          : nullptr;
+  if (!ConstScale || ConstScale->getZExtValue() != 1) {
+    VPValue *Scale = IV->getOperand(1);
----------------
fhahn wrote:

Done , thanks!

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