Repository: llvm/llvm-project
File: llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (C++)

Code before:
  TestTransformFn Fn;
};

} // namespace

// Wrappers to adapt different function signatures to uniform TestTransformFn.
template <void (*Fn)(VPlan &)>
static void wrapTransform(VPlan &P, const TargetLibraryInfo *) {
  Fn(P);
}

template <bool (*Fn)(VPlan &)>
static void wrapTransform(VPlan &P, const TargetLibraryInfo *) {
  Fn(P);
}

static const TestTransformEntry TestTransforms[] = {
    {"simplify-recipes", wrapTransform<VPlanTransforms::simplifyRecipes>},
    {"remove-dead-recipes", wrapTransform<VPlanTransforms::removeDeadRecipes>},
    {"simplify-blends", wrapTransform<simplifyBlends>},
    {"merge-blocks", wrapTransform<mergeBlocksIntoPredecessors>},
    {"licm", wrapTransform<licm>},
    {"cse", wrapTransform<VPlanTransforms::cse>},
    {"optimize", wrapTransform<VPlanTransforms::optimize>},
    {"widen-from-metadata", VPlanTransforms::widenFromMetadata},
};

static const TestTransformEntry *getTestTransform(StringRef Name) {
  for (const auto &E : TestTransforms)
    if (E.Name == Name)
      return &E;
  return nullptr;
}

void VPlanTransforms::runTestTransforms(VPlan &Plan, StringRef Pipeline,
                                        const TargetLibraryInfo *TLI) {
  SmallVector<StringRef> Passes;
  Pipeline.split(Passes, ',');

  for (StringRef PassName : Passes) {
    PassName = PassName.trim();
    if (PassName == "print") {
      dbgs() << "VPlan:\n" << Plan << "\n";
      continue;
    }
    const TestTransformEntry *E = getTestTransform(PassName);
    if (!E)
      report_fatal_error("Unknown VPlan test transform: " + PassName);
    E->Fn(Plan, TLI);
  }
}

Review comment (refactor):
Maybe we can extract this out into a lightweight VPlanPassManager, and change runTransform use-sites to use it?

Code after:
    auto *S = new VPWidenStoreRecipe(
        *SI, StoreGroup->getAddr(), Res, nullptr, /*Consecutive=*/true,
        /*Reverse=*/false, {}, StoreGroup->getDebugLoc());
    S->insertBefore(StoreGroup);
    StoreGroup->eraseFromParent();
  }

  // Adjust induction to reflect that the transformed plan only processes one
  // original iteration.
  auto *CanIV = VectorLoop->getCanonicalIV();
  auto *Inc = cast<VPInstruction>(CanIV->getBackedgeValue());
  VPBuilder PHBuilder(Plan.getVectorPreheader());

  VPValue *UF = &Plan.getUF();
  if (VFToOptimize->isScalable()) {
    VPValue *VScale = PHBuilder.createElementCount(
        VectorLoop->getCanonicalIVType(), ElementCount::getScalable(1));
    VPValue *VScaleUF = PHBuilder.createOverflowingOp(
        Instruction::Mul, {VScale, UF}, {true, false});
    Inc->setOperand(1, VScaleUF);
    Plan.getVF().replaceAllUsesWith(VScale);
  } else {
    Inc->setOperand(1, UF);
    Plan.getVF().replaceAllUsesWith(
        Plan.getConstantInt(CanIV->getScalarType(), 1));
  }
  removeDeadRecipes(Plan);
  assert(none_of(*VectorLoop->getEntryBasicBlock(),
                 IsaPred<VPVectorPointerRecipe>) &&
         "All VPVectorPointerRecipes should have been removed");
  return NewPlan;
}

/// Add branch weight metadata, if the \p Plan's middle block is terminated by a
/// BranchOnCond recipe.
void VPlanTransforms::addBranchWeightToMiddleTerminator(
    VPlan &Plan, ElementCount VF, std::optional<unsigned> VScaleForTuning) {
  VPBasicBlock *MiddleVPBB = Plan.getMiddleBlock();
  auto *MiddleTerm =
      dyn_cast_or_null<VPInstruction>(MiddleVPBB->getTerminator());
  // Only add branch metadata if there is a (conditional) terminator.
  if (!MiddleTerm)
    return;

  assert(MiddleTerm->getOpcode() == VPInstruction::BranchOnCond &&
         "must have a BranchOnCond");
  // Assume that `TripCount % VectorStep ` is equally distributed.
  unsigned VectorStep = Plan.getConcreteUF() * VF.getKnownMinValue();
  if (VF.isScalable() && VScaleForTuning.has_value())
    VectorStep *= *VScaleForTuning;
  assert(VectorStep > 0 && "trip count should not be zero");