[VP] Merge ExpandVP pass into PreISelIntrinsicLowering (#101652)

Similar to #97727; avoid an extra pass over the entire IR by performing
the lowering as part of the pre-isel-intrinsic-lowering pass.

1 files changed, 35 insertions, 103 deletions

diff --git a/llvm/lib/CodeGen/ExpandVectorPredication.cpp b/llvm/lib/CodeGen/ExpandVectorPredication.cpp
index 97c6ee4..5ffdbcd 100644
--- a/llvm/lib/CodeGen/ExpandVectorPredication.cpp
+++ b/llvm/lib/CodeGen/ExpandVectorPredication.cpp
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This pass implements IR expansion for vector predication intrinsics, allowing
+// This file implements IR expansion for vector predication intrinsics, allowing
 // targets to enable vector predication until just before codegen.
 //
 //===----------------------------------------------------------------------===//
@@ -16,7 +16,6 @@
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/VectorUtils.h"
-#include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
@@ -24,8 +23,6 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
-#include "llvm/InitializePasses.h"
-#include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
@@ -137,7 +134,6 @@ namespace {
 
 // Expansion pass state at function scope.
 struct CachingVPExpander {
-  Function &F;
   const TargetTransformInfo &TTI;
 
   /// \returns A (fixed length) vector with ascending integer indices
@@ -207,10 +203,10 @@ struct CachingVPExpander {
   bool UsingTTIOverrides;
 
 public:
-  CachingVPExpander(Function &F, const TargetTransformInfo &TTI)
-      : F(F), TTI(TTI), UsingTTIOverrides(anyExpandVPOverridesSet()) {}
+  CachingVPExpander(const TargetTransformInfo &TTI)
+      : TTI(TTI), UsingTTIOverrides(anyExpandVPOverridesSet()) {}
 
-  bool expandVectorPredication();
+  bool expandVectorPredication(VPIntrinsic &VPI);
 };
 
 //// CachingVPExpander {
@@ -571,7 +567,7 @@ CachingVPExpander::expandPredicationInMemoryIntrinsic(IRBuilder<> &Builder,
                                                       VPIntrinsic &VPI) {
   assert(VPI.canIgnoreVectorLengthParam());
 
-  const auto &DL = F.getDataLayout();
+  const auto &DL = VPI.getDataLayout();
 
   Value *MaskParam = VPI.getMaskParam();
   Value *PtrParam = VPI.getMemoryPointerParam();
@@ -775,15 +771,6 @@ Value *CachingVPExpander::expandPredication(VPIntrinsic &VPI) {
 
 //// } CachingVPExpander
 
-struct TransformJob {
-  VPIntrinsic *PI;
-  TargetTransformInfo::VPLegalization Strategy;
-  TransformJob(VPIntrinsic *PI, TargetTransformInfo::VPLegalization InitStrat)
-      : PI(PI), Strategy(InitStrat) {}
-
-  bool isDone() const { return Strategy.shouldDoNothing(); }
-};
-
 void sanitizeStrategy(VPIntrinsic &VPI, VPLegalization &LegalizeStrat) {
   // Operations with speculatable lanes do not strictly need predication.
   if (maySpeculateLanes(VPI)) {
@@ -821,98 +808,43 @@ CachingVPExpander::getVPLegalizationStrategy(const VPIntrinsic &VPI) const {
 }
 
 /// Expand llvm.vp.* intrinsics as requested by \p TTI.
-bool CachingVPExpander::expandVectorPredication() {
-  SmallVector<TransformJob, 16> Worklist;
-
-  // Collect all VPIntrinsics that need expansion and determine their expansion
-  // strategy.
-  for (auto &I : instructions(F)) {
-    auto *VPI = dyn_cast<VPIntrinsic>(&I);
-    if (!VPI)
-      continue;
-    auto VPStrat = getVPLegalizationStrategy(*VPI);
-    sanitizeStrategy(*VPI, VPStrat);
-    if (!VPStrat.shouldDoNothing())
-      Worklist.emplace_back(VPI, VPStrat);
-  }
-  if (Worklist.empty())
-    return false;
+bool CachingVPExpander::expandVectorPredication(VPIntrinsic &VPI) {
+  auto Strategy = getVPLegalizationStrategy(VPI);
+  sanitizeStrategy(VPI, Strategy);
 
-  // Transform all VPIntrinsics on the worklist.
-  LLVM_DEBUG(dbgs() << "\n:::: Transforming " << Worklist.size()
-                    << " instructions ::::\n");
-  for (TransformJob Job : Worklist) {
-    // Transform the EVL parameter.
-    switch (Job.Strategy.EVLParamStrategy) {
-    case VPLegalization::Legal:
-      break;
-    case VPLegalization::Discard:
-      discardEVLParameter(*Job.PI);
-      break;
-    case VPLegalization::Convert:
-      if (foldEVLIntoMask(*Job.PI))
-        ++NumFoldedVL;
-      break;
-    }
-    Job.Strategy.EVLParamStrategy = VPLegalization::Legal;
+  // Transform the EVL parameter.
+  switch (Strategy.EVLParamStrategy) {
+  case VPLegalization::Legal:
+    break;
+  case VPLegalization::Discard:
+    discardEVLParameter(VPI);
+    break;
+  case VPLegalization::Convert:
+    if (foldEVLIntoMask(VPI))
+      ++NumFoldedVL;
+    break;
+  }
 
-    // Replace with a non-predicated operation.
-    switch (Job.Strategy.OpStrategy) {
-    case VPLegalization::Legal:
-      break;
-    case VPLegalization::Discard:
-      llvm_unreachable("Invalid strategy for operators.");
-    case VPLegalization::Convert:
-      expandPredication(*Job.PI);
+  // Replace with a non-predicated operation.
+  switch (Strategy.OpStrategy) {
+  case VPLegalization::Legal:
+    break;
+  case VPLegalization::Discard:
+    llvm_unreachable("Invalid strategy for operators.");
+  case VPLegalization::Convert:
+    if (Value *V = expandPredication(VPI); V != &VPI) {
       ++NumLoweredVPOps;
-      break;
+      // Return true if and only if the intrinsic was actually removed.
+      return true;
     }
-    Job.Strategy.OpStrategy = VPLegalization::Legal;
-
-    assert(Job.isDone() && "incomplete transformation");
+    break;
   }
 
-  return true;
+  return false;
 }
-class ExpandVectorPredication : public FunctionPass {
-public:
-  static char ID;
-  ExpandVectorPredication() : FunctionPass(ID) {
-    initializeExpandVectorPredicationPass(*PassRegistry::getPassRegistry());
-  }
-
-  bool runOnFunction(Function &F) override {
-    const auto *TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
-    CachingVPExpander VPExpander(F, *TTI);
-    return VPExpander.expandVectorPredication();
-  }
-
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<TargetTransformInfoWrapperPass>();
-    AU.setPreservesCFG();
-  }
-};
 } // namespace
 
-char ExpandVectorPredication::ID;
-INITIALIZE_PASS_BEGIN(ExpandVectorPredication, "expandvp",
-                      "Expand vector predication intrinsics", false, false)
-INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
-INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
-INITIALIZE_PASS_END(ExpandVectorPredication, "expandvp",
-                    "Expand vector predication intrinsics", false, false)
-
-FunctionPass *llvm::createExpandVectorPredicationPass() {
-  return new ExpandVectorPredication();
-}
-
-PreservedAnalyses
-ExpandVectorPredicationPass::run(Function &F, FunctionAnalysisManager &AM) {
-  const auto &TTI = AM.getResult<TargetIRAnalysis>(F);
-  CachingVPExpander VPExpander(F, TTI);
-  if (!VPExpander.expandVectorPredication())
-    return PreservedAnalyses::all();
-  PreservedAnalyses PA;
-  PA.preserveSet<CFGAnalyses>();
-  return PA;
+bool llvm::expandVectorPredicationIntrinsic(VPIntrinsic &VPI,
+                                            const TargetTransformInfo &TTI) {
+  return CachingVPExpander(TTI).expandVectorPredication(VPI);
 }