diff --git a/llvm/include/llvm/Analysis/ML/InliningAdvisor.h b/llvm/include/llvm/Analysis/ML/InliningAdvisor.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Analysis/ML/InliningAdvisor.h @@ -0,0 +1,154 @@ +//===- InliningAdvisor.h - ML infrastructure for inliner --------*- C++ -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +#ifndef LLVM_ML_INLINERML_H_ +#define LLVM_ML_INLINERML_H_ + +#include +#include + +#include "llvm/IR/PassManager.h" + +namespace llvm { +class CallBase; +class Function; +class Module; +class PreservedAnalyses; + +enum class MLMode : int { Invalid, Rel, Dev }; + +/// Capture state between an inlining decision having had been made, and +/// its impact being observable. When collecting model training data, this +/// allows recording features/decisions/partial reward data sets. +class PendingInliningRecord { +public: + PendingInliningRecord(PendingInliningRecord &&) = delete; + PendingInliningRecord(const PendingInliningRecord &) = delete; + virtual ~PendingInliningRecord() = default; + + /// Call after inlining succeeded, and did not result in deleting the callee. + void recordInlining() { + recordInlining(/*CalleeWasDeleted*/ false, /*SiteWasInlined*/ true); + } + + /// Call after inlining succeeded, and resulted in deleting the callee. + void recordInliningWithCalleeDeleted() { + recordInlining(/*CalleeWasDeleted*/ true, /*SiteWasInlined*/ true); + } + + /// Call after the decision for a call site was to not inline. + void recordUnsucessfulInlining() { + recordInlining(/*CalleeWasDeleted*/ false, /*SiteWasInlined*/ false); + } + +protected: + PendingInliningRecord() = default; + + /// Update internal state if an inlining was successful (or callee was + /// deleted). In addition, if we are in 'dev' mode, and collecting training + /// logs, compute partial reward and log the inlining event, i.e. the features + /// observed before inlining, the inlining decision, as well as the partial + /// reward. + virtual void recordInlining(bool CalleeWasDeleted, bool SiteWasInlined) = 0; +}; + +/// The InliningAdvisor extracts features from a call site and surrounding +/// context (e.g. module) and interfaces with a ML model executor (normally AOT, +/// may optionally be JIT or interpreted, if experimentation is desired) to +/// evaluate whether inlining should take place. +class InliningAdvisor { +public: + InliningAdvisor(InliningAdvisor &&) = delete; + virtual ~InliningAdvisor() = default; + + /// Modify the \p AlternativeRecommendation with the advice from the trained + /// model about whether to inline \p CB or not. CB is assumed to be a direct + /// call. If logging (or for debug), capture the call site features provided + /// at this point in the returned PendingInliningRecord. + /// \p Mandatory indicates the inlining must happen because of always inline + /// attributes, and in that case, the advisor won't interfere, but it will + /// want to update its internal reprensetation based on whether the inlining + /// succeeds. + /// Returns nullptr if the inlining is illegal, or a PendingInliningRecord to + /// keep track of the state between when the inlining decision is made and + /// after the inlining happens. The user is responsible to call the + /// appropriate record* method on the PendingInliningRecord, to inform the + /// advisor about what happened and allow it to observe partial rewards, log, + /// etc. + virtual std::unique_ptr + getInliningAdvice(CallBase &CB, bool &AlternativeRecommendation, + bool Mandatory) = 0; + + virtual void OnPassEntry() = 0; + virtual void OnPassExit() = 0; + + virtual void OnAllInliningCompleted() = 0; + + static std::unique_ptr + create(Module &, ModuleAnalysisManager &, MLMode Mode); + + /// Return true if the advisor is performing inference. This means we don't + /// need to evaluate using the manual heuristic whether to inline or not, + /// since the result would be ignored. + bool isPerformingInference() const { return IsPerformingInference; } + +protected: + InliningAdvisor() = default; + void setIsPerformingInferenceOn(); + +private: + bool IsPerformingInference = false; +}; + +/// The InliningAdvisorAnalysis is a module pass, albeit, currently, the inliner +/// is a SCC pass. This is because the InliningAdvisor needs to capture state +/// right before inlining commences over a module. +class InliningAdvisorAnalysis + : public AnalysisInfoMixin { +public: + InliningAdvisorAnalysis(MLMode Mode) : Mode(Mode) {} + struct Result { + Result(std::unique_ptr Adv) : Advisor(std::move(Adv)) { + assert(Advisor && "Creating an InliningAdvisorAnalysis without a valid " + "InliningAdvisor is not supported"); + } + bool invalidate(Module &, const PreservedAnalyses &, + ModuleAnalysisManager::Invalidator &) { + // InliningAdvisor must be preserved across analysis invalidations. + return false; + } + InliningAdvisor *get() const { return Advisor.get(); } + + private: + std::unique_ptr Advisor; + }; + + Result run(Module &M, ModuleAnalysisManager &MAM); + static AnalysisKey Key; + +private: + static bool isModeSupported(MLMode Mode); + MLMode Mode = MLMode::Invalid; +}; + +/// We defer deleting functions to after the inlining for a whole module has +/// finished. This allows us to reliably use function pointers as unique +/// identifiers, as an efficient implementation detail of the InliningAdvisor. +/// Otherwise, it is possible the memory allocator re-allocate Function objects +/// at the same address of a deleted Function; and Functions are potentially +/// created during the function passes called after each SCC inlining (e.g. +/// argument promotion does that). +class InliningAdvisorCleanupPass + : public PassInfoMixin { +public: + InliningAdvisorCleanupPass() = default; + PreservedAnalyses run(Module &, ModuleAnalysisManager &); +}; + +} // namespace llvm +#endif // LLVM_ML_INLINERML_H_ diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt --- a/llvm/lib/Analysis/CMakeLists.txt +++ b/llvm/lib/Analysis/CMakeLists.txt @@ -40,6 +40,7 @@ IVUsers.cpp IndirectCallPromotionAnalysis.cpp InlineCost.cpp + InliningAdvisorAnalysis.cpp InstCount.cpp InstructionPrecedenceTracking.cpp InstructionSimplify.cpp diff --git a/llvm/lib/Analysis/InliningAdvisorAnalysis.cpp b/llvm/lib/Analysis/InliningAdvisorAnalysis.cpp new file mode 100644 --- /dev/null +++ b/llvm/lib/Analysis/InliningAdvisorAnalysis.cpp @@ -0,0 +1,51 @@ +//===- InliningAdvisorAnalysis.cpp - analysis pass implementation +//----------===// +// +// The LLVM Compiler Infrastructure +// +// This file is distributed under the University of Illinois Open Source +// License. See LICENSE.TXT for details. +// +//===----------------------------------------------------------------------===// +// +// This file implements InliningAdvisorAnalysis and InliningAdvisorCleanupPass. +// +//===----------------------------------------------------------------------===// + +#include "llvm/Analysis/ML/InliningAdvisor.h" + +namespace llvm { + +std::unique_ptr +createAdvisor(Module &M, ModuleAnalysisManager &MAM, MLMode Mode) { + switch (Mode) { + case MLMode::Invalid: + llvm_unreachable("The MLMode::Invalid case should have been handled in the " + "pass manager"); + return nullptr; + case MLMode::Dev: + case MLMode::Rel: +#if LLVM_HAVE_TF_API || LLVM_HAVE_TF_AOT + return InliningAdvisor::create(M, MAM, Mode); +#endif + break; + } + M.getContext().emitError( + "Could not setup Inlining Advisor for the requested mode and/or options"); + return nullptr; +} + +AnalysisKey InliningAdvisorAnalysis::Key; +InliningAdvisorAnalysis::Result +InliningAdvisorAnalysis::run(Module &M, ModuleAnalysisManager &MAM) { + return Result(createAdvisor(M, MAM, Mode)); +} + +PreservedAnalyses InliningAdvisorCleanupPass::run(Module &M, + ModuleAnalysisManager &MAM) { + auto *Advisor = MAM.getCachedResult(M); + if (Advisor) + Advisor->get()->OnAllInliningCompleted(); + return PreservedAnalyses::all(); +} +} // namespace llvm diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp --- a/llvm/lib/Passes/PassBuilder.cpp +++ b/llvm/lib/Passes/PassBuilder.cpp @@ -39,6 +39,7 @@ #include "llvm/Analysis/LoopCacheAnalysis.h" #include "llvm/Analysis/LoopInfo.h" #include "llvm/Analysis/LoopNestAnalysis.h" +#include "llvm/Analysis/ML/InliningAdvisor.h" #include "llvm/Analysis/MemoryDependenceAnalysis.h" #include "llvm/Analysis/MemorySSA.h" #include "llvm/Analysis/ModuleSummaryAnalysis.h" @@ -215,6 +216,16 @@ "enable-npm-gvn-hoist", cl::init(false), cl::Hidden, cl::desc("Enable the GVN hoisting pass for the new PM (default = off)")); +static cl::opt EnableMLInliner( + "enable-ml-inliner", cl::init(MLMode::Invalid), cl::Hidden, + cl::desc("Enable ML policy for inliner. Currently trained for -Oz only"), + cl::values(clEnumValN(MLMode::Invalid, "disabled", + "Heuristics-based inliner version."), + clEnumValN(MLMode::Dev, "dev", + "Use development mode (runtime-loadable model)."), + clEnumValN(MLMode::Rel, "rel", + "Use release mode (AOT-compiled model)."))); + static cl::opt EnableGVNSink( "enable-npm-gvn-sink", cl::init(false), cl::Hidden, cl::desc("Enable the GVN hoisting pass for the new PM (default = off)")); @@ -715,8 +726,12 @@ if (Phase == ThinLTOPhase::PreLink && PGOOpt && PGOOpt->Action == PGOOptions::SampleUse) IP.HotCallSiteThreshold = 0; - MainCGPipeline.addPass(InlinerPass(IP)); + if (EnableMLInliner != MLMode::Invalid) { + MPM.addPass(RequireAnalysisPass()); + } + + MainCGPipeline.addPass(InlinerPass(IP)); if (AttributorRun & AttributorRunOption::CGSCC) MainCGPipeline.addPass(AttributorCGSCCPass()); @@ -752,6 +767,8 @@ MPM.addPass( createModuleToPostOrderCGSCCPassAdaptor(createDevirtSCCRepeatedPass( std::move(MainCGPipeline), MaxDevirtIterations))); + if (EnableMLInliner != MLMode::Invalid) + MPM.addPass(InliningAdvisorCleanupPass()); return MPM; } diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def --- a/llvm/lib/Passes/PassRegistry.def +++ b/llvm/lib/Passes/PassRegistry.def @@ -27,6 +27,7 @@ MODULE_ANALYSIS("verify", VerifierAnalysis()) MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC)) MODULE_ANALYSIS("asan-globals-md", ASanGlobalsMetadataAnalysis()) +MODULE_ANALYSIS("inlining-advisor", InliningAdvisorAnalysis(EnableMLInliner)) #ifndef MODULE_ALIAS_ANALYSIS #define MODULE_ALIAS_ANALYSIS(NAME, CREATE_PASS) \ diff --git a/llvm/lib/Transforms/IPO/Inliner.cpp b/llvm/lib/Transforms/IPO/Inliner.cpp --- a/llvm/lib/Transforms/IPO/Inliner.cpp +++ b/llvm/lib/Transforms/IPO/Inliner.cpp @@ -30,6 +30,7 @@ #include "llvm/Analysis/CallGraph.h" #include "llvm/Analysis/InlineCost.h" #include "llvm/Analysis/LazyCallGraph.h" +#include "llvm/Analysis/ML/InliningAdvisor.h" #include "llvm/Analysis/OptimizationRemarkEmitter.h" #include "llvm/Analysis/ProfileSummaryInfo.h" #include "llvm/Analysis/TargetLibraryInfo.h" @@ -890,6 +891,24 @@ assert(InitialC.size() > 0 && "Cannot handle an empty SCC!"); Module &M = *InitialC.begin()->getFunction().getParent(); ProfileSummaryInfo *PSI = MAM.getCachedResult(M); + auto *IAA = MAM.getCachedResult(M); + InliningAdvisor *Advisor = IAA ? IAA->get() : nullptr; + if (Advisor) + Advisor->OnPassEntry(); + + // Avoid subtle bugs due to alternative exits from this method - if we have + // an advisor, ensure it is always informed when we're done with a scc. + class AdvisorExitCapture final { + InliningAdvisor *const Advisor; + + public: + AdvisorExitCapture(InliningAdvisor *A) : Advisor(A) {} + ~AdvisorExitCapture() { + if (Advisor) + Advisor->OnPassExit(); + } + }; + AdvisorExitCapture Capturer(Advisor); if (!ImportedFunctionsStats && InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No) { @@ -1059,17 +1078,42 @@ continue; } - Optional OIC = shouldInline(*CS, GetInlineCost, ORE); - // Check whether we want to inline this callsite. - if (!OIC.hasValue()) { - setInlineRemark(*CS, "deferred"); + auto TrivialDecision = llvm::getAttributeBasedInliningDecision( + *CS, CS->getCalledFunction(), FAM.getResult(Callee), + GetTLI); + + if (Advisor && + ((TrivialDecision.hasValue() && !TrivialDecision->isSuccess()) || + &Callee == &F)) continue; - } - if (!OIC.getValue()) { - // shouldInline() call returned a negative inline cost that explains - // why this callsite should not be inlined. - setInlineRemark(*CS, inlineCostStr(*OIC)); + const bool Mandatory = + TrivialDecision.hasValue() && TrivialDecision->isSuccess(); + + Optional OIC; + // TODO(mtrofin): shouldInline replicates the already calculated + // TrivialDecision part when we don't do inference. Refactor to avoid. + if (!Advisor || !Advisor->isPerformingInference()) + OIC = shouldInline(*CS, GetInlineCost, ORE); + assert(!Mandatory || (OIC.hasValue() && OIC.getValue())); + bool ShouldInline = Mandatory || (OIC.hasValue() && OIC.getValue()); + + std::unique_ptr PendingRecord; + if (Advisor) + PendingRecord = + Advisor->getInliningAdvice(*CS, ShouldInline, Mandatory); + + if (!ShouldInline) { + // Check whether we want to inline this callsite. + if (!OIC.hasValue()) { + setInlineRemark(*CS, "deferred"); + } else if (!OIC.getValue()) { + // shouldInline() call returned a negative inline cost that explains + // why this callsite should not be inlined. + setInlineRemark(*CS, inlineCostStr(*OIC)); + } + if (PendingRecord) + PendingRecord->recordUnsucessfulInlining(); continue; } @@ -1088,14 +1132,17 @@ InlineResult IR = InlineFunction(*CS, IFI); if (!IR.isSuccess()) { - setInlineRemark(*CS, std::string(IR.getFailureReason()) + "; " + - inlineCostStr(*OIC)); + setInlineRemark( + *CS, std::string(IR.getFailureReason()) + "; " + + (OIC.hasValue() ? inlineCostStr(*OIC) : "ML Advisor")); ORE.emit([&]() { return OptimizationRemarkMissed(DEBUG_TYPE, "NotInlined", DLoc, Block) << NV("Callee", &Callee) << " will not be inlined into " << NV("Caller", &F) << ": " << NV("Reason", IR.getFailureReason()); }); + if (PendingRecord) + PendingRecord->recordUnsucessfulInlining(); continue; } DidInline = true; @@ -1103,7 +1150,10 @@ ++NumInlined; - emitInlinedInto(ORE, DLoc, Block, Callee, F, *OIC); + // TODO(mtrofin): OIC may not have value if Advisor decided against + // inlining. We should still emit a remark. + if (OIC.hasValue()) + emitInlinedInto(ORE, DLoc, Block, Callee, F, *OIC); // Add any new callsites to defined functions to the worklist. if (!IFI.InlinedCallSites.empty()) { @@ -1136,6 +1186,7 @@ // 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). @@ -1155,8 +1206,14 @@ assert(find(DeadFunctions, &Callee) == DeadFunctions.end() && "Cannot put cause a function to become dead twice!"); DeadFunctions.push_back(&Callee); + CalleeWasDeleted = true; } } + if (PendingRecord) + if (CalleeWasDeleted) + PendingRecord->recordInliningWithCalleeDeleted(); + else + PendingRecord->recordInlining(); } // Back the call index up by one to put us in a good position to go around @@ -1247,7 +1304,19 @@ UR.InvalidatedRefSCCs.insert(&DeadRC); // And delete the actual function from the module. - M.getFunctionList().erase(DeadF); + // If we use the Advisor, it uses Function pointers to index various + // maps, e.g. memoization. Function cleanup passes like argument promotion + // create new functions. It is possible for a new function to be allocated + // at the address of a deleted function. + // We could index using names, but that's inefficient. Alternatively, + // we let the Advisor free the functions. + if (Advisor) { + DeadF->getBasicBlockList().clear(); + M.getFunctionList().remove(DeadF); + } else { + M.getFunctionList().erase(DeadF); + } + ++NumDeleted; } diff --git a/llvm/test/Transforms/Inline/inlining-advisor-default.ll b/llvm/test/Transforms/Inline/inlining-advisor-default.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/Inline/inlining-advisor-default.ll @@ -0,0 +1,9 @@ +; Check that, in the absence of dependencies, we emit an error message when +; trying to use ML-driven inlining. +; +; RUN: not opt -passes=scc-oz-module-inliner -enable-ml-inliner=dev -S < %s 2>&1 | FileCheck %s +; RUN: not opt -passes=scc-oz-module-inliner -enable-ml-inliner=rel -S < %s 2>&1 | FileCheck %s + +declare i64 @f1() + +; CHECK: Could not setup Inlining Advisor for the requested mode and/or options \ No newline at end of file