[ArgPromotion] Remove all the getters and ReplaceCallSite

[psamolysov]

Summary

AARGetter is an abstraction over a source of the AAResults introduced
to support the legacy pass manager as well as the modern one. Since the
Argument Promotion pass doesn't support the legacy pass manager anymore,
the abstraction is not required and AAResults may be used directly.

The instance of the FunctionAnalysisManager is passed through the
functions to get all the required analyses just wherever they are
required and do not use the awkward getter callbacks.

The ReplaceCallSite parameter was required for the legacy pass manager
only and isn't used anymore, so the parameter has been eliminated.

[nikic]

The disadvantage is that AAResult gets computed unconditionally, while previously it was only computed if the function was a viable promotion candidate.

Maybe we want to pass through FunctionAnalysisManager to avoid the awkward callbacks, while keeping the current positions of the lookup?

[psamolysov]

Good catch, thank you. I think the ReplicaCallSite parameter can be removed too because we pass only None for it.

[psamolysov]

Eliminate all the getters as well as the ReplaceCallSite parameter.

[psamolysov]

Remove #include of unused headers.

[aeubanks]

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

44

some of these includes should still be here, like CGSCCPassManager, LazyCallGraph, and Module

[psamolysov]

CGSCCPassManager and LazyCallGraph are included in the class' header. About Module I'm not sure, CLion marks this header as unused, I also see no Module in the code.

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

 // would be best handled when and if LLVM begins supporting multiple return
 // values from functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 
 #include "llvm/ADT/DepthFirstIterator.h"
-#include "llvm/ADT/None.h"
-#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
-#include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/CallGraph.h"
-#include "llvm/Analysis/CallGraphSCCPass.h"
-#include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/MemoryLocation.h"
-#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Metadata.h"
-#include "llvm/IR/Module.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
-#include "llvm/InitializePasses.h"
-#include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
-#include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
if (OrigOffset != 0) {
     Ptr = IRB.CreateGEP(IRB.getInt8Ty(), Ptr, IRB.getInt(OrigOffset));
   }
   return IRB.CreateBitCast(Ptr, ResElemTy->getPointerTo(AddrSpace));
 }
 
 /// DoPromotion - This method actually performs the promotion of the specified
 /// arguments, and returns the new function.  At this point, we know that it's
 /// safe to do so.
-static Function *doPromotion(
-    Function *F, function_ref<DominatorTree &(Function &F)> DTGetter,
-    function_ref<AssumptionCache *(Function &F)> ACGetter,
-    const DenseMap<Argument *, SmallVector<OffsetAndArgPart, 4>> &ArgsToPromote,
-    Optional<function_ref<void(CallBase &OldCS, CallBase &NewCS)>>
-        ReplaceCallSite) {
+static Function *
+doPromotion(Function *F, FunctionAnalysisManager &FAM,
+            const DenseMap<Argument *, SmallVector<OffsetAndArgPart, 4>>
+                &ArgsToPromote) {
   // Start by computing a new prototype for the function, which is the same as
   // the old function, but has modified arguments.
   FunctionType *FTy = F->getFunctionType();
   std::vector<Type *> Params;
 
   // Attribute - Keep track of the parameter attributes for the arguments
   // that we are *not* promoting. For the ones that we do promote, the parameter
   // attributes are lost
doPromotion(Function *F, FunctionAnalysisManager &FAM,
   NF->setAttributes(AttributeList::get(F->getContext(), PAL.getFnAttrs(),
                                        PAL.getRetAttrs(), ArgAttrVec));
   AttributeFuncs::updateMinLegalVectorWidthAttr(*NF, LargestVectorWidth);
   ArgAttrVec.clear();
 
   F->getParent()->getFunctionList().insert(F->getIterator(), NF);
   NF->takeName(F);
 
-  // Loop over all of the callers of the function, transforming the call sites
-  // to pass in the loaded pointers.
-  //
+  // Loop over all the callers of the function, transforming the call sites to
+  // pass in the loaded pointers.
   SmallVector<Value *, 16> Args;
   const DataLayout &DL = F->getParent()->getDataLayout();
   while (!F->use_empty()) {
     CallBase &CB = cast<CallBase>(*F->user_back());
     assert(CB.getCalledFunction() == F);
     const AttributeList &CallPAL = CB.getAttributes();
     IRBuilder<NoFolder> IRB(&CB);
 
NewCS->setAttributes(AttributeList::get(F->getContext(),
                                             CallPAL.getRetAttrs(), ArgAttrVec));
     NewCS->copyMetadata(CB, {LLVMContext::MD_prof, LLVMContext::MD_dbg});
     Args.clear();
     ArgAttrVec.clear();
 
     AttributeFuncs::updateMinLegalVectorWidthAttr(*CB.getCaller(),
                                                   LargestVectorWidth);
 
-    // Update the callgraph to know that the callsite has been transformed.
-    if (ReplaceCallSite)
-      (*ReplaceCallSite)(CB, *NewCS);
-
     if (!CB.use_empty()) {
       CB.replaceAllUsesWith(NewCS);
       NewCS->takeName(&CB);
     }
 
     // Finally, remove the old call from the program, reducing the use-count of
     // F.
     CB.eraseFromParent();
doPromotion(Function *F, FunctionAnalysisManager &FAM,
 
   LLVM_DEBUG(dbgs() << "ARG PROMOTION: " << Allocas.size()
                     << " alloca(s) are promotable by Mem2Reg\n");
 
   if (!Allocas.empty()) {
     // And we are able to call the `promoteMemoryToRegister()` function.
     // Our earlier checks have ensured that PromoteMemToReg() will
     // succeed.
-    PromoteMemToReg(Allocas, DTGetter(*NF), ACGetter(*NF));
+    auto &DT = FAM.getResult<DominatorTreeAnalysis>(*NF);
+    auto &AC = FAM.getResult<AssumptionAnalysis>(*NF);
+    PromoteMemToReg(Allocas, DT, &AC);
   }
 
   return NF;
 }
 
 /// Return true if we can prove that all callees pass in a valid pointer for the
 /// specified function argument.
 static bool allCallersPassValidPointerForArgument(Argument *Arg,
return all_of(F.uses(), [&](const Use &U) {
     return TTI.areTypesABICompatible(Caller, Callee, Types);
   });
 }
 
 /// PromoteArguments - This method checks the specified function to see if there
 /// are any promotable arguments and if it is safe to promote the function (for
 /// example, all callers are direct).  If safe to promote some arguments, it
 /// calls the DoPromotion method.
-static Function *
-promoteArguments(Function *F, function_ref<AAResults &(Function &F)> AARGetter,
-                 function_ref<DominatorTree &(Function &F)> DTGetter,
-                 function_ref<AssumptionCache *(Function &F)> ACGetter,
-                 unsigned MaxElements,
-                 Optional<function_ref<void(CallBase &OldCS, CallBase &NewCS)>>
-                     ReplaceCallSite,
-                 const TargetTransformInfo &TTI, bool IsRecursive) {
+static Function *promoteArguments(Function *F, FunctionAnalysisManager &FAM,
+                                  unsigned MaxElements, bool IsRecursive) {
   // Don't perform argument promotion for naked functions; otherwise we can end
   // up removing parameters that are seemingly 'not used' as they are referred
   // to in the assembly.
   if (F->hasFnAttribute(Attribute::Naked))
     return nullptr;
 
   // Make sure that it is local to this module.
   if (!F->hasLocalLinkage())
static Function *promoteArguments(Function *F, FunctionAnalysisManager &FAM,
 
   // Can't change signature of musttail caller
   // FIXME: Support promoting whole chain of musttail functions
   for (BasicBlock &BB : *F)
     if (BB.getTerminatingMustTailCall())
       return nullptr;
 
   const DataLayout &DL = F->getParent()->getDataLayout();
-
-  AAResults &AAR = AARGetter(*F);
+  auto &AAR = FAM.getResult<AAManager>(*F);
+  const auto &TTI = FAM.getResult<TargetIRAnalysis>(*F);
 
   // Check to see which arguments are promotable.  If an argument is promotable,
   // add it to ArgsToPromote.
   DenseMap<Argument *, SmallVector<OffsetAndArgPart, 4>> ArgsToPromote;
   for (Argument *PtrArg : PointerArgs) {
     // Replace sret attribute with noalias. This reduces register pressure by
     // avoiding a register copy.
     if (PtrArg->hasStructRetAttr()) {
if (findArgParts(PtrArg, DL, AAR, MaxElements, IsRecursive, ArgParts)) {
       }
     }
   }
 
   // No promotable pointer arguments.
   if (ArgsToPromote.empty())
     return nullptr;
 
-  return doPromotion(F, DTGetter, ACGetter, ArgsToPromote, ReplaceCallSite);
+  return doPromotion(F, FAM, ArgsToPromote);
 }
 
 PreservedAnalyses ArgumentPromotionPass::run(LazyCallGraph::SCC &C,
                                              CGSCCAnalysisManager &AM,
                                              LazyCallGraph &CG,
                                              CGSCCUpdateResult &UR) {
   bool Changed = false, LocalChange;
 
   // Iterate until we stop promoting from this SCC.
   do {
     LocalChange = false;
 
     FunctionAnalysisManager &FAM =
         AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
     bool IsRecursive = C.size() > 1;
     for (LazyCallGraph::Node &N : C) {
       Function &OldF = N.getFunction();
-
-      // FIXME: This lambda must only be used with this function. We should
-      // skip the lambda and just get the AA results directly.
-      auto AARGetter = [&](Function &F) -> AAResults & {
-        assert(&F == &OldF && "Called with an unexpected function!");
-        return FAM.getResult<AAManager>(F);
-      };
-
-      auto DTGetter = [&](Function &F) -> DominatorTree & {
-        assert(&F != &OldF && "Called with the obsolete function!");
-        return FAM.getResult<DominatorTreeAnalysis>(F);
-      };
-
-      auto ACGetter = [&](Function &F) -> AssumptionCache * {
-        assert(&F != &OldF && "Called with the obsolete function!");
-        return &FAM.getResult<AssumptionAnalysis>(F);
-      };
-
-      const auto &TTI = FAM.getResult<TargetIRAnalysis>(OldF);
-      Function *NewF = promoteArguments(&OldF, AARGetter, DTGetter, ACGetter,
-                                        MaxElements, None, TTI, IsRecursive);
+      Function *NewF = promoteArguments(&OldF, FAM, MaxElements, IsRecursive);
       if (!NewF)
         continue;
       LocalChange = true;
 
       // Directly substitute the functions in the call graph. Note that this
       // requires the old function to be completely dead and completely
       // replaced by the new function. It does no call graph updates, it merely
       // swaps out the particular function mapped to a particular node in the