[llvm-reduce] optimize extractFromModule functions

[dwightguth]

Summary

The extractBasicBlocksFromModule, extractInstrFromModule, and other
similar functions previously performed very poorly when the number of
such elements in the program to reduce was very high. Previously, we
were creating the set which caches elements to keep by looping through
all elements in the module and adding them to the set. However, since
std::set is an ordered set, this introduces a massive amount of
rebalancing if the order of elements in the program and the order of
their pointers in memory are not the same.

The solution is straightforward: first put all the elements to be kept
in a vector, then use the constructor for std::set which takes a pair of
iterators over a collection. This constructor is optimized to avoid
doing unnecessary work when initializing large sets.

Also in this change, we pass BBsToKeep set to functions
replaceBranchTerminator and removeUninterestingBBsFromSwitch as a const
reference rather than passing it by value. This ought to prevent the
need to copy the collection each time these functions are called, which
is expensive if the collection is large.

[dwightguth]

removed stray change made by mistake

[aeubanks]

does std::unordered_set help?

[dwightguth]

I tried changing it to an unordered set without any of the other changes originally, but it still had the same problem of a lot of time spent inserting into the hash table. I eventually abandoned that approach because my intuition was that a hash table where the keys are addresses in memory would probably end up poorly balanced a lot of the time.

If you want I can try out the combination of pushing first into a vector and also creating a hash set from the vector, but I didn't go that route because I had already largely addressed the performance bottleneck in this function with the diff as you see it now, and it didn't seem worth the risk that it might be badly balanced in some cases when the majority of the execution time was now taking place elsewhere.

[aeubanks]

adding some comments saying that this is for performance reasons would be good

[dwightguth]

add comments

[dwightguth]

I addressed your suggestion that we add comments. I do not have commit access yet; can you please commit it if you're happy with the diff?

llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp

 //===- ReduceArguments.cpp - Specialized Delta Pass -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 #include <vector>
 
 using namespace llvm;
 
 /// Goes over OldF calls and replaces them with a call to NewF
 static void replaceFunctionCalls(Function &OldF, Function &NewF,
                                  const std::set<int> &ArgIndexesToKeep) {
   const auto &Users = OldF.users();
-  for (auto I = Users.begin(), E = Users.end(); I != E; )
+  for (auto I = Users.begin(), E = Users.end(); I != E;)
     if (auto *CI = dyn_cast<CallInst>(*I++)) {
       // Skip uses in call instructions where OldF isn't the called function
       // (e.g. if OldF is an argument of the call).
       if (CI->getCalledFunction() != &OldF)
         continue;
       SmallVector<Value *, 8> Args;
       for (auto ArgI = CI->arg_begin(), E = CI->arg_end(); ArgI != E; ++ArgI)
         if (ArgIndexesToKeep.count(ArgI - CI->arg_begin()))

 /// fixed.
 static bool shouldRemoveArguments(const Function &F) {
   return !F.arg_empty() && !F.isIntrinsic();
 }
 
 /// Removes out-of-chunk arguments from functions, and modifies their calls
 /// accordingly. It also removes allocations of out-of-chunk arguments.
 static void extractArgumentsFromModule(Oracle &O, Module &Program) {
-  std::set<Argument *> ArgsToKeep;
+  std::vector<Argument *> InitArgsToKeep;
   std::vector<Function *> Funcs;
   // Get inside-chunk arguments, as well as their parent function
   for (auto &F : Program)
     if (shouldRemoveArguments(F)) {
       Funcs.push_back(&F);
       for (auto &A : F.args())
         if (O.shouldKeep())
-          ArgsToKeep.insert(&A);
+          InitArgsToKeep.push_back(&A);
     }
 
+  std::set<Argument *> ArgsToKeep(InitArgsToKeep.begin(), InitArgsToKeep.end());
+
   for (auto *F : Funcs) {
     ValueToValueMapTy VMap;
     std::vector<WeakVH> InstToDelete;
     for (auto &A : F->args())
       if (!ArgsToKeep.count(&A)) {
         // By adding undesired arguments to the VMap, CloneFunction will remove
         // them from the resulting Function
         VMap[&A] = UndefValue::get(A.getType());

llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp

 //===- ReduceArguments.cpp - Specialized Delta Pass -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 #include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
 #include <vector>
 
 using namespace llvm;
 
 /// Replaces BB Terminator with one that only contains Chunk BBs
 static void replaceBranchTerminator(BasicBlock &BB,
-                                    std::set<BasicBlock *> BBsToKeep) {
+                                    const std::set<BasicBlock *> &BBsToKeep) {
   auto *Term = BB.getTerminator();
   std::vector<BasicBlock *> ChunkSucessors;
   for (auto *Succ : successors(&BB))
     if (BBsToKeep.count(Succ))
       ChunkSucessors.push_back(Succ);
 
   // BB only references Chunk BBs
   if (ChunkSucessors.size() == Term->getNumSuccessors())
if (Address) {
     for (auto *Dest : ChunkSucessors)
       NewIndBI->addDestination(Dest);
   }
 }
 
 /// Removes uninteresting BBs from switch, if the default case ends up being
 /// uninteresting, the switch is replaced with a void return (since it has to be
 /// replace with something)
-static void removeUninterestingBBsFromSwitch(SwitchInst &SwInst,
-                                             std::set<BasicBlock *> BBsToKeep) {
+static void
+removeUninterestingBBsFromSwitch(SwitchInst &SwInst,
+                                 const std::set<BasicBlock *> &BBsToKeep) {
   if (!BBsToKeep.count(SwInst.getDefaultDest())) {
     auto *FnRetTy = SwInst.getParent()->getParent()->getReturnType();
     ReturnInst::Create(SwInst.getContext(),
                        FnRetTy->isVoidTy() ? nullptr : UndefValue::get(FnRetTy),
                        SwInst.getParent());
     SwInst.eraseFromParent();
   } else
     for (int I = 0, E = SwInst.getNumCases(); I != E; ++I) {
       auto Case = SwInst.case_begin() + I;
       if (!BBsToKeep.count(Case->getCaseSuccessor())) {
         SwInst.removeCase(Case);
         --I;
         --E;
       }
     }
 }
 
 /// Removes out-of-chunk arguments from functions, and modifies their calls
 /// accordingly. It also removes allocations of out-of-chunk arguments.
 static void extractBasicBlocksFromModule(Oracle &O, Module &Program) {
-  std::set<BasicBlock *> BBsToKeep;
+  std::vector<BasicBlock *> InitBBsToKeep;
 
   for (auto &F : Program)
     for (auto &BB : F)
       if (O.shouldKeep())
-        BBsToKeep.insert(&BB);
+        InitBBsToKeep.push_back(&BB);
+
+  std::set<BasicBlock *> BBsToKeep(InitBBsToKeep.begin(), InitBBsToKeep.end());
 
   std::vector<BasicBlock *> BBsToDelete;
   for (auto &F : Program)
     for (auto &BB : F) {
       if (!BBsToKeep.count(&BB)) {
         BBsToDelete.push_back(&BB);
         // Remove out-of-chunk BB from successor phi nodes
         for (auto *Succ : successors(&BB))

llvm/tools/llvm-reduce/deltas/ReduceGlobalVars.cpp

 //===- ReduceGlobalVars.cpp - Specialized Delta Pass ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a function which calls the Generic Delta pass in order
 // to reduce Global Variables in the provided Module.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ReduceGlobalVars.h"
 #include "llvm/IR/Constants.h"
 #include <set>
 
 using namespace llvm;
 
 /// Removes all the GVs that aren't inside the desired Chunks.
 static void extractGVsFromModule(Oracle &O, Module &Program) {
   // Get GVs inside desired chunks
-  std::set<GlobalVariable *> GVsToKeep;
+  std::vector<GlobalVariable *> InitGVsToKeep;
   for (auto &GV : Program.globals())
     if (O.shouldKeep())
-      GVsToKeep.insert(&GV);
+      InitGVsToKeep.push_back(&GV);
+
+  std::set<GlobalVariable *> GVsToKeep(InitGVsToKeep.begin(),
+                                       InitGVsToKeep.end());
 
   // Delete out-of-chunk GVs and their uses
   std::vector<GlobalVariable *> ToRemove;
   std::vector<WeakVH> InstToRemove;
   for (auto &GV : Program.globals())
     if (!GVsToKeep.count(&GV)) {
       for (auto *U : GV.users())
         if (auto *Inst = dyn_cast<Instruction>(U))

llvm/tools/llvm-reduce/deltas/ReduceInstructions.cpp

 //===- ReduceArguments.cpp - Specialized Delta Pass -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a function which calls the Generic Delta pass in order
 // to reduce uninteresting Arguments from defined functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ReduceInstructions.h"
 
 using namespace llvm;
 
 /// Removes out-of-chunk arguments from functions, and modifies their calls
 /// accordingly. It also removes allocations of out-of-chunk arguments.
 static void extractInstrFromModule(Oracle &O, Module &Program) {
-  std::set<Instruction *> InstToKeep;
+  std::vector<Instruction *> InitInstToKeep;
 
   for (auto &F : Program)
     for (auto &BB : F) {
       // Removing the terminator would make the block invalid. Only iterate over
       // instructions before the terminator.
-      InstToKeep.insert(BB.getTerminator());
+      InitInstToKeep.push_back(BB.getTerminator());
       for (auto &Inst : make_range(BB.begin(), std::prev(BB.end())))
         if (O.shouldKeep())
-          InstToKeep.insert(&Inst);
+          InitInstToKeep.push_back(&Inst);
     }
 
+  std::set<Instruction *> InstToKeep(InitInstToKeep.begin(),
+                                     InitInstToKeep.end());
+
   std::vector<Instruction *> InstToDelete;
   for (auto &F : Program)
     for (auto &BB : F)
       for (auto &Inst : BB)
         if (!InstToKeep.count(&Inst)) {
           Inst.replaceAllUsesWith(UndefValue::get(Inst.getType()));
           InstToDelete.push_back(&Inst);
         }