[Transforms] Merge function attributes within InlineFunction (NFC)

[kazu]

Summary

In the past, we've had a bug resulting in a compiler crash after
forgetting to merge function attributes (D105729).

This patch teaches InlineFunction to merge function attributes. This
way, we minimize the "time" when the IR is valid, but the function
attributes are not.

[davidxl]

Why not adding another parameter to InlineFunction method to specify whether attribute Merge is done?

[kazu]

Taught InlineFunction to take an additional parameter.

[kazu]

In D134117#3797926, @davidxl wrote:

Why not adding another parameter to InlineFunction method to specify whether attribute Merge is done?

Thank you for the review. I've updated the patch. Please take a look.

One concern is that I have to pass several default parameters before I can specify the last parameter. Perhaps I could reorder the parameters to reduce the number default parameters to pass, but that should be left as a follow-up patch.

[davidxl]

lgtm (The InlineFunction body is getting really big. Probably worth breaking it up in the future).

[kazu]

Thanks for the review!

In D134117#3797982, @davidxl wrote:

lgtm (The InlineFunction body is getting really big. Probably worth breaking it up in the future).

Totally agreed.

llvm/include/llvm/Transforms/Utils/Cloning.h

 /// *inlined* code to minimize the actual inserted code, it must not delete
 /// code in the caller as users of this routine may have pointers to
 /// instructions in the caller that need to remain stable.
 ///
 /// If ForwardVarArgsTo is passed, inlining a function with varargs is allowed
 /// and all varargs at the callsite will be passed to any calls to
 /// ForwardVarArgsTo. The caller of InlineFunction has to make sure any varargs
 /// are only used by ForwardVarArgsTo.
+///
+/// The callee's function attributes are merged into the callers' if
+/// MergeAttributes is set to true.
 InlineResult InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
                             AAResults *CalleeAAR = nullptr,
                             bool InsertLifetime = true,
-                            Function *ForwardVarArgsTo = nullptr);
+                            Function *ForwardVarArgsTo = nullptr,
+                            bool MergeAttributes = false);
 
 /// Clones a loop \p OrigLoop.  Returns the loop and the blocks in \p
 /// Blocks.
 ///
 /// Updates LoopInfo and DominatorTree assuming the loop is dominated by block
 /// \p LoopDomBB.  Insert the new blocks before block specified in \p Before.
 /// Note: Only innermost loops are supported.
 Loop *cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,

llvm/lib/Transforms/IPO/AlwaysInliner.cpp

if (!F.isDeclaration() && isInlineViable(F).isSuccess()) {
         BasicBlock *Block = CB->getParent();
 
         InlineFunctionInfo IFI(
             /*cg=*/nullptr, GetAssumptionCache, &PSI,
             &FAM.getResult<BlockFrequencyAnalysis>(*Caller),
             &FAM.getResult<BlockFrequencyAnalysis>(F));
 
         InlineResult Res = InlineFunction(
-            *CB, IFI, &FAM.getResult<AAManager>(F), InsertLifetime);
+            *CB, IFI, &FAM.getResult<AAManager>(F), InsertLifetime,
+            /*ForwardVarArgsTo=*/nullptr,
+            /*MergeAttributes=*/true);
         if (!Res.isSuccess()) {
           ORE.emit([&]() {
             return OptimizationRemarkMissed(DEBUG_TYPE, "NotInlined", DLoc,
                                             Block)
                    << "'" << ore::NV("Callee", &F) << "' is not inlined into '"
                    << ore::NV("Caller", Caller)
                    << "': " << ore::NV("Reason", Res.getFailureReason());
           });
           continue;
         }
 
         emitInlinedIntoBasedOnCost(
             ORE, DLoc, Block, F, *Caller,
             InlineCost::getAlways("always inline attribute"),
             /*ForProfileContext=*/false, DEBUG_TYPE);
 
-        // Merge the attributes based on the inlining.
-        AttributeFuncs::mergeAttributesForInlining(*Caller, F);
-
         Changed = true;
       }
 
       if (F.hasFnAttribute(Attribute::AlwaysInline)) {
         // Remember to try and delete this function afterward. This both avoids
         // re-walking the rest of the module and avoids dealing with any
         // iterator invalidation issues while deleting functions.
         InlinedFunctions.push_back(&F);

llvm/lib/Transforms/IPO/Inliner.cpp

static InlineResult inlineCallIfPossible(
     ImportedFunctionsInliningStatistics &ImportedFunctionsStats) {
   Function *Callee = CB.getCalledFunction();
   Function *Caller = CB.getCaller();
 
   AAResults &AAR = AARGetter(*Callee);
 
   // Try to inline the function.  Get the list of static allocas that were
   // inlined.
-  InlineResult IR = InlineFunction(CB, IFI, &AAR, InsertLifetime);
+  InlineResult IR = InlineFunction(CB, IFI, &AAR, InsertLifetime,
+                                   /*ForwardVarArgsTo=*/nullptr,
+                                   /*MergeAttributes=*/true);
   if (!IR.isSuccess())
     return IR;
 
   if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
     ImportedFunctionsStats.recordInline(*Caller, *Callee);
 
-  AttributeFuncs::mergeAttributesForInlining(*Caller, *Callee);
-
   if (!DisableInlinedAllocaMerging)
     mergeInlinedArrayAllocas(Caller, IFI, InlinedArrayAllocas, InlineHistory);
 
   return IR; // success
 }
 
 /// Return true if the specified inline history ID
 /// indicates an inline history that includes the specified function.
for (; I < (int)Calls.size() && Calls[I].first->getCaller() == &F; ++I) {
       // Setup the data structure used to plumb customization into the
       // `InlineFunction` routine.
       InlineFunctionInfo IFI(
           /*cg=*/nullptr, GetAssumptionCache, PSI,
           &FAM.getResult<BlockFrequencyAnalysis>(*(CB->getCaller())),
           &FAM.getResult<BlockFrequencyAnalysis>(Callee));
 
       InlineResult IR =
-          InlineFunction(*CB, IFI, &FAM.getResult<AAManager>(*CB->getCaller()));
+          InlineFunction(*CB, IFI, &FAM.getResult<AAManager>(*CB->getCaller()),
+                         /*InsertLifetime=*/true,
+                         /*ForwardVarArgsTo=*/nullptr,
+                         /*MergeAttributes=*/true);
       if (!IR.isSuccess()) {
         Advice->recordUnsuccessfulInlining(IR);
         continue;
       }
 
       DidInline = true;
       InlinedCallees.insert(&Callee);
       ++NumInlined;
for (; I < (int)Calls.size() && Calls[I].first->getCaller() == &F; ++I) {
                     itostr(CBCostMult * IntraSCCCostMultiplier));
                 ICB->addFnAttr(NewCBCostMult);
               }
             }
           }
         }
       }
 
-      // Merge the attributes based on the inlining.
-      AttributeFuncs::mergeAttributesForInlining(F, Callee);
-
       // For local functions or discardable functions without comdats, check
       // whether this makes the callee trivially dead. In that case, we can drop
       // the body of the function eagerly which may reduce the number of callers
       // of other functions to one, changing inline cost thresholds. Non-local
       // discardable functions with comdats are checked later on.
       bool CalleeWasDeleted = false;
       if (Callee.isDiscardableIfUnused() && Callee.hasZeroLiveUses() &&
           !CG.isLibFunction(Callee)) {

llvm/lib/Transforms/IPO/ModuleInliner.cpp

while (!Calls->empty()) {
     // Setup the data structure used to plumb customization into the
     // `InlineFunction` routine.
     InlineFunctionInfo IFI(
         /*cg=*/nullptr, GetAssumptionCache, PSI,
         &FAM.getResult<BlockFrequencyAnalysis>(*(CB->getCaller())),
         &FAM.getResult<BlockFrequencyAnalysis>(Callee));
 
     InlineResult IR =
-        InlineFunction(*CB, IFI, &FAM.getResult<AAManager>(*CB->getCaller()));
+        InlineFunction(*CB, IFI, &FAM.getResult<AAManager>(*CB->getCaller()),
+                       /*InsertLifetime=*/true,
+                       /*ForwardVarArgsTo=*/nullptr,
+                       /*MergeAttributes=*/true);
     if (!IR.isSuccess()) {
       Advice->recordUnsuccessfulInlining(IR);
       continue;
     }
 
     Changed = true;
     ++NumInlined;
 
if (!IFI.InlinedCallSites.empty()) {
             NewCallee = ICB->getCalledFunction();
         }
         if (NewCallee)
           if (!NewCallee->isDeclaration())
             Calls->push({ICB, NewHistoryID});
       }
     }
 
-    // Merge the attributes based on the inlining.
-    AttributeFuncs::mergeAttributesForInlining(F, Callee);
-
     // For local functions, check whether this makes the callee trivially
     // dead. In that case, we can drop the body of the function eagerly
     // which may reduce the number of callers of other functions to one,
     // changing inline cost thresholds.
     bool CalleeWasDeleted = false;
     if (Callee.hasLocalLinkage()) {
       // To check this we also need to nuke any dead constant uses (perhaps
       // made dead by this operation on other functions).

llvm/lib/Transforms/IPO/SampleProfile.cpp

if (Cost.isNever()) {
     return false;
   }
 
   if (!Cost)
     return false;
 
   InlineFunctionInfo IFI(nullptr, GetAC);
   IFI.UpdateProfile = false;
-  if (!InlineFunction(CB, IFI).isSuccess())
+  InlineResult IR = InlineFunction(CB, IFI,
+                                   /*CalleeAAR=*/nullptr,
+                                   /*InsertLifetime=*/true,
+                                   /*ForwardVarArgsTo=*/nullptr,
+                                   /*MergeAttributes=*/true);
+  if (!IR.isSuccess())
     return false;
 
-  // Merge the attributes based on the inlining.
-  AttributeFuncs::mergeAttributesForInlining(*BB->getParent(),
-                                             *CalledFunction);
-
   // The call to InlineFunction erases I, so we can't pass it here.
   emitInlinedIntoBasedOnCost(*ORE, DLoc, BB, *CalledFunction, *BB->getParent(),
                              Cost, true, getAnnotatedRemarkPassName());
 
   // Now populate the list of newly exposed call sites.
   if (InlinedCallSites) {
     InlinedCallSites->clear();
     for (auto &I : IFI.InlinedCallSites)

llvm/lib/Transforms/Utils/InlineFunction.cpp

 ///
 /// Note that this only does one level of inlining.  For example, if the
 /// instruction 'call B' is inlined, and 'B' calls 'C', then the call to 'C' now
 /// exists in the instruction stream.  Similarly this will inline a recursive
 /// function by one level.
 llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
                                         AAResults *CalleeAAR,
                                         bool InsertLifetime,
-                                        Function *ForwardVarArgsTo) {
+                                        Function *ForwardVarArgsTo,
+                                        bool MergeAttributes) {
   assert(CB.getParent() && CB.getFunction() && "Instruction not in function!");
 
   // FIXME: we don't inline callbr yet.
   if (isa<CallBrInst>(CB))
     return InlineResult::failure("We don't inline callbr yet.");
 
   // If IFI has any state in it, zap it before we fill it in.
   IFI.reset();
if (!CB.use_empty()) {
         CB.replaceAllUsesWith(R->getReturnValue());
     }
     // Since we are now done with the Call/Invoke, we can delete it.
     CB.eraseFromParent();
 
     // Since we are now done with the return instruction, delete it also.
     Returns[0]->eraseFromParent();
 
+    if (MergeAttributes)
+      AttributeFuncs::mergeAttributesForInlining(*Caller, *CalledFunc);
+
     // We are now done with the inlining.
     return InlineResult::success();
   }
 
   // Otherwise, we have the normal case, of more than one block to inline or
   // multiple return sites.
 
   // We want to clone the entire callee function into the hole between the
AssumptionCache *AC =
         IFI.GetAssumptionCache ? &IFI.GetAssumptionCache(*Caller) : nullptr;
     auto &DL = Caller->getParent()->getDataLayout();
     if (Value *V = simplifyInstruction(PHI, {DL, nullptr, nullptr, AC})) {
       PHI->replaceAllUsesWith(V);
       PHI->eraseFromParent();
     }
   }
 
+  if (MergeAttributes)
+    AttributeFuncs::mergeAttributesForInlining(*Caller, *CalledFunc);
+
   return InlineResult::success();
 }