[llvm] [SandboxVec][BottomUpVec] Fix ownership of Legality (PR #143018)

[via llvm-commits]

https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/143018

>From 05d3b75da08f3315e0291557a2dfd0bfd9333808 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Fri, 7 Mar 2025 15:00:09 -0800
Subject: [PATCH] [SandboxVec][BottomUpVec] Fix ownership of Legality

Fix the ownership of `Legality` member variable of BottomUpVec.
It should get created in runOnFunction() and get destroyed when the function
returns.
---
 .../SandboxVectorizer/Passes/BottomUpVec.h    |  5 ++--
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  | 28 +++++++++++--------
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
index af0e07de4f51f..c65e2cd45ab48 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
@@ -34,7 +34,6 @@ namespace llvm::sandboxir {
 /// transaction, depending on the cost.
 class BottomUpVec final : public RegionPass {
   bool Change = false;
-  std::unique_ptr<LegalityAnalysis> Legality;
   /// The original instructions that are potentially dead after vectorization.
   DenseSet<Instruction *> DeadInstrCandidates;
   /// Maps scalars to vectors.
@@ -88,12 +87,12 @@ class BottomUpVec final : public RegionPass {
   /// Recursively try to vectorize \p Bndl and its operands. This populates the
   /// `Actions` vector.
   Action *vectorizeRec(ArrayRef<Value *> Bndl, ArrayRef<Value *> UserBndl,
-                       unsigned Depth);
+                       unsigned Depth, LegalityAnalysis &Legality);
   /// Generate vector instructions based on `Actions` and return the last vector
   /// created.
   Value *emitVectors();
   /// Entry point for vectorization starting from \p Seeds.
-  bool tryVectorize(ArrayRef<Value *> Seeds);
+  bool tryVectorize(ArrayRef<Value *> Seeds, LegalityAnalysis &Legality);
 
 public:
   BottomUpVec() : RegionPass("bottom-up-vec") {}
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 9781d4cf146ef..86dbd2171a560 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -278,13 +278,14 @@ void BottomUpVec::collectPotentiallyDeadInstrs(ArrayRef<Value *> Bndl) {
 }
 
 Action *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl,
-                                  ArrayRef<Value *> UserBndl, unsigned Depth) {
+                                  ArrayRef<Value *> UserBndl, unsigned Depth,
+                                  LegalityAnalysis &Legality) {
   bool StopForDebug =
       DebugBndlCnt++ >= StopBundle && StopBundle != StopBundleDisabled;
   LLVM_DEBUG(dbgs() << DEBUG_PREFIX << "canVectorize() Bundle:\n";
              VecUtils::dump(Bndl));
-  const auto &LegalityRes = StopForDebug ? Legality->getForcedPackForDebugging()
-                                         : Legality->canVectorize(Bndl);
+  const auto &LegalityRes = StopForDebug ? Legality.getForcedPackForDebugging()
+                                         : Legality.canVectorize(Bndl);
   LLVM_DEBUG(dbgs() << DEBUG_PREFIX << "Legality: " << LegalityRes << "\n");
   auto ActionPtr =
       std::make_unique<Action>(&LegalityRes, Bndl, UserBndl, Depth);
@@ -297,14 +298,16 @@ Action *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl,
       break;
     case Instruction::Opcode::Store: {
       // Don't recurse towards the pointer operand.
-      Action *OpA = vectorizeRec(getOperand(Bndl, 0), Bndl, Depth + 1);
+      Action *OpA =
+          vectorizeRec(getOperand(Bndl, 0), Bndl, Depth + 1, Legality);
       Operands.push_back(OpA);
       break;
     }
     default:
       // Visit all operands.
       for (auto OpIdx : seq<unsigned>(I->getNumOperands())) {
-        Action *OpA = vectorizeRec(getOperand(Bndl, OpIdx), Bndl, Depth + 1);
+        Action *OpA =
+            vectorizeRec(getOperand(Bndl, OpIdx), Bndl, Depth + 1, Legality);
         Operands.push_back(OpA);
       }
       break;
@@ -476,16 +479,17 @@ Value *BottomUpVec::emitVectors() {
   return NewVec;
 }
 
-bool BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl) {
+bool BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl,
+                               LegalityAnalysis &Legality) {
   Change = false;
   if (LLVM_UNLIKELY(BottomUpInvocationCnt++ >= StopAt &&
                     StopAt != StopAtDisabled))
     return false;
   DeadInstrCandidates.clear();
-  Legality->clear();
+  Legality.clear();
   Actions.clear();
   DebugBndlCnt = 0;
-  vectorizeRec(Bndl, {}, /*Depth=*/0);
+  vectorizeRec(Bndl, {}, /*Depth=*/0, Legality);
   LLVM_DEBUG(dbgs() << DEBUG_PREFIX << "BottomUpVec: Vectorization Actions:\n";
              Actions.dump());
   emitVectors();
@@ -498,16 +502,16 @@ bool BottomUpVec::runOnRegion(Region &Rgn, const Analyses &A) {
   assert(SeedSlice.size() >= 2 && "Bad slice!");
   Function &F = *SeedSlice[0]->getParent()->getParent();
   IMaps = std::make_unique<InstrMaps>();
-  Legality = std::make_unique<LegalityAnalysis>(
-      A.getAA(), A.getScalarEvolution(), F.getParent()->getDataLayout(),
-      F.getContext(), *IMaps);
+  LegalityAnalysis Legality(A.getAA(), A.getScalarEvolution(),
+                            F.getParent()->getDataLayout(), F.getContext(),
+                            *IMaps);
 
   // TODO: Refactor to remove the unnecessary copy to SeedSliceVals.
   SmallVector<Value *> SeedSliceVals(SeedSlice.begin(), SeedSlice.end());
   // Try to vectorize starting from the seed slice. The returned value
   // is true if we found vectorizable code and generated some vector
   // code for it. It does not mean that the code is profitable.
-  return tryVectorize(SeedSliceVals);
+  return tryVectorize(SeedSliceVals, Legality);
 }
 
 } // namespace sandboxir