[llvm] [SROA] Use SetVector for PromotableAllocas (PR #105809)

[Bartłomiej Chmiel via llvm-commits]

https://github.com/b-chmiel updated https://github.com/llvm/llvm-project/pull/105809

>From 54bbd117cad05baa21e15c909553b003a2224df2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bart=C5=82omiej=20Chmiel?= <bchmiel at antmicro.com>
Date: Thu, 22 Aug 2024 14:56:21 +0200
Subject: [PATCH 1/2] [SROA] Use SmallPtrSet for PromotableAllocas
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Optimize SROA pass for large number of allocas by
speeding-up PromotableAllocas erase operation. The optimization
involves using SmallPtrSet which proves to be efficient since
PromotableAllocas is used only for manipulating unique pointers.

Signed-off-by: Bartłomiej Chmiel <bchmiel at antmicro.com>
---
 llvm/lib/Transforms/Scalar/SROA.cpp | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 2310cb3a7decbb..4ead8240c17e77 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -198,7 +198,7 @@ class SROA {
   SmallSetVector<AllocaInst *, 16> PostPromotionWorklist;
 
   /// A collection of alloca instructions we can directly promote.
-  std::vector<AllocaInst *> PromotableAllocas;
+  SmallPtrSet<AllocaInst *, 16> PromotableAllocas;
 
   /// A worklist of PHIs to speculate prior to promoting allocas.
   ///
@@ -4799,9 +4799,8 @@ bool SROA::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
 
   // Finally, don't try to promote any allocas that new require re-splitting.
   // They have already been added to the worklist above.
-  llvm::erase_if(PromotableAllocas, [&](AllocaInst *AI) {
-    return ResplitPromotableAllocas.count(AI);
-  });
+  for (auto *RPA : ResplitPromotableAllocas)
+    PromotableAllocas.erase(RPA);
 
   return true;
 }
@@ -4963,7 +4962,7 @@ AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
     }
     if (PHIUsers.empty() && SelectUsers.empty()) {
       // Promote the alloca.
-      PromotableAllocas.push_back(NewAI);
+      PromotableAllocas.insert(NewAI);
     } else {
       // If we have either PHIs or Selects to speculate, add them to those
       // worklists and re-queue the new alloca so that we promote in on the
@@ -5598,7 +5597,9 @@ bool SROA::promoteAllocas(Function &F) {
     LLVM_DEBUG(dbgs() << "Not promoting allocas with mem2reg!\n");
   } else {
     LLVM_DEBUG(dbgs() << "Promoting allocas with mem2reg...\n");
-    PromoteMemToReg(PromotableAllocas, DTU->getDomTree(), AC);
+    PromoteMemToReg(
+        std::vector(PromotableAllocas.begin(), PromotableAllocas.end()),
+        DTU->getDomTree(), AC);
   }
 
   PromotableAllocas.clear();
@@ -5615,7 +5616,7 @@ std::pair<bool /*Changed*/, bool /*CFGChanged*/> SROA::runSROA(Function &F) {
     if (AllocaInst *AI = dyn_cast<AllocaInst>(I)) {
       if (DL.getTypeAllocSize(AI->getAllocatedType()).isScalable() &&
           isAllocaPromotable(AI))
-        PromotableAllocas.push_back(AI);
+        PromotableAllocas.insert(AI);
       else
         Worklist.insert(AI);
     }
@@ -5639,10 +5640,10 @@ std::pair<bool /*Changed*/, bool /*CFGChanged*/> SROA::runSROA(Function &F) {
       // Remove the deleted allocas from various lists so that we don't try to
       // continue processing them.
       if (!DeletedAllocas.empty()) {
-        auto IsInSet = [&](AllocaInst *AI) { return DeletedAllocas.count(AI); };
-        Worklist.remove_if(IsInSet);
-        PostPromotionWorklist.remove_if(IsInSet);
-        llvm::erase_if(PromotableAllocas, IsInSet);
+        Worklist.set_subtract(DeletedAllocas);
+        PostPromotionWorklist.set_subtract(DeletedAllocas);
+        for (auto *DA : DeletedAllocas)
+          PromotableAllocas.erase(DA);
         DeletedAllocas.clear();
       }
     }

>From b0e3b7ebcf87e8fa50574342306f3083f54da001 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bart=C5=82omiej=20Chmiel?= <bchmiel at antmicro.com>
Date: Mon, 26 Aug 2024 10:29:39 +0200
Subject: [PATCH 2/2] [SROA] Use SetVector for PromotableAllocas
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Bartłomiej Chmiel <bchmiel at antmicro.com>
---
 llvm/lib/Transforms/Scalar/SROA.cpp | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 4ead8240c17e77..d0186da1bc5e22 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -198,7 +198,9 @@ class SROA {
   SmallSetVector<AllocaInst *, 16> PostPromotionWorklist;
 
   /// A collection of alloca instructions we can directly promote.
-  SmallPtrSet<AllocaInst *, 16> PromotableAllocas;
+  SetVector<AllocaInst *, SmallVector<AllocaInst *>,
+            SmallPtrSet<AllocaInst *, 16>, 16>
+      PromotableAllocas;
 
   /// A worklist of PHIs to speculate prior to promoting allocas.
   ///
@@ -4799,8 +4801,7 @@ bool SROA::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
 
   // Finally, don't try to promote any allocas that new require re-splitting.
   // They have already been added to the worklist above.
-  for (auto *RPA : ResplitPromotableAllocas)
-    PromotableAllocas.erase(RPA);
+  PromotableAllocas.set_subtract(ResplitPromotableAllocas);
 
   return true;
 }
@@ -5597,9 +5598,7 @@ bool SROA::promoteAllocas(Function &F) {
     LLVM_DEBUG(dbgs() << "Not promoting allocas with mem2reg!\n");
   } else {
     LLVM_DEBUG(dbgs() << "Promoting allocas with mem2reg...\n");
-    PromoteMemToReg(
-        std::vector(PromotableAllocas.begin(), PromotableAllocas.end()),
-        DTU->getDomTree(), AC);
+    PromoteMemToReg(PromotableAllocas.getArrayRef(), DTU->getDomTree(), AC);
   }
 
   PromotableAllocas.clear();
@@ -5642,8 +5641,7 @@ std::pair<bool /*Changed*/, bool /*CFGChanged*/> SROA::runSROA(Function &F) {
       if (!DeletedAllocas.empty()) {
         Worklist.set_subtract(DeletedAllocas);
         PostPromotionWorklist.set_subtract(DeletedAllocas);
-        for (auto *DA : DeletedAllocas)
-          PromotableAllocas.erase(DA);
+        PromotableAllocas.set_subtract(DeletedAllocas);
         DeletedAllocas.clear();
       }
     }