diff --git a/llvm/include/llvm/Analysis/InlineAdvisor.h b/llvm/include/llvm/Analysis/InlineAdvisor.h --- a/llvm/include/llvm/Analysis/InlineAdvisor.h +++ b/llvm/include/llvm/Analysis/InlineAdvisor.h @@ -357,7 +357,8 @@ }; std::unique_ptr -getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM); +getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM, + std::function GetDefaultAdvice); std::unique_ptr getDevelopmentModeAdvisor(Module &M, ModuleAnalysisManager &MAM, diff --git a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h --- a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h +++ b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h @@ -134,6 +134,7 @@ extern const char *const DecisionName; extern const TensorSpec InlineDecisionSpec; extern const char *const DefaultDecisionName; +extern const TensorSpec DefaultDecisionSpec; extern const char *const RewardName; using InlineFeatures = std::vector; diff --git a/llvm/include/llvm/Analysis/MLInlineAdvisor.h b/llvm/include/llvm/Analysis/MLInlineAdvisor.h --- a/llvm/include/llvm/Analysis/MLInlineAdvisor.h +++ b/llvm/include/llvm/Analysis/MLInlineAdvisor.h @@ -28,7 +28,8 @@ class MLInlineAdvisor : public InlineAdvisor { public: MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr ModelRunner); + std::unique_ptr ModelRunner, + std::function GetDefaultAdvice); virtual ~MLInlineAdvisor() = default; @@ -63,6 +64,7 @@ unsigned getInitialFunctionLevel(const Function &F) const; std::unique_ptr ModelRunner; + std::function GetDefaultAdvice; private: int64_t getModuleIRSize() const; diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp --- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp +++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp @@ -165,7 +165,6 @@ bool isLogging() const { return !!Logger; } std::unique_ptr getMandatoryAdviceImpl(CallBase &CB) override; - std::function GetDefaultAdvice; const bool IsDoingInference; std::unique_ptr Logger; @@ -280,7 +279,7 @@ append_range(FT, MUTR->extraOutputsForLoggingSpecs()); DefaultDecisionPos = FT.size(); - FT.push_back(TensorSpec::createSpec(DefaultDecisionName, {1})); + FT.push_back(DefaultDecisionSpec); DecisionPos = FT.size(); FT.push_back(InlineDecisionSpec); @@ -331,8 +330,7 @@ std::unique_ptr ModelRunner, std::function GetDefaultAdvice, std::unique_ptr Logger) - : MLInlineAdvisor(M, MAM, std::move(ModelRunner)), - GetDefaultAdvice(GetDefaultAdvice), + : MLInlineAdvisor(M, MAM, std::move(ModelRunner), GetDefaultAdvice), IsDoingInference(isa(getModelRunner())), Logger(std::move(Logger)), InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0), diff --git a/llvm/lib/Analysis/InlineAdvisor.cpp b/llvm/lib/Analysis/InlineAdvisor.cpp --- a/llvm/lib/Analysis/InlineAdvisor.cpp +++ b/llvm/lib/Analysis/InlineAdvisor.cpp @@ -208,6 +208,10 @@ Advisor.reset(DA.Factory(M, FAM, Params, IC)); return !!Advisor; } + auto GetDefaultAdvice = [&FAM, Params](CallBase &CB) { + auto OIC = getDefaultInlineAdvice(CB, FAM, Params); + return OIC.has_value(); + }; switch (Mode) { case InliningAdvisorMode::Default: LLVM_DEBUG(dbgs() << "Using default inliner heuristic.\n"); @@ -223,16 +227,12 @@ case InliningAdvisorMode::Development: #ifdef LLVM_HAVE_TFLITE LLVM_DEBUG(dbgs() << "Using development-mode inliner policy.\n"); - Advisor = - llvm::getDevelopmentModeAdvisor(M, MAM, [&FAM, Params](CallBase &CB) { - auto OIC = getDefaultInlineAdvice(CB, FAM, Params); - return OIC.has_value(); - }); + Advisor = llvm::getDevelopmentModeAdvisor(M, MAM, GetDefaultAdvice); #endif break; case InliningAdvisorMode::Release: LLVM_DEBUG(dbgs() << "Using release-mode inliner policy.\n"); - Advisor = llvm::getReleaseModeAdvisor(M, MAM); + Advisor = llvm::getReleaseModeAdvisor(M, MAM, GetDefaultAdvice); break; } diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp --- a/llvm/lib/Analysis/MLInlineAdvisor.cpp +++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp @@ -38,6 +38,12 @@ "Base file path for the interactive mode. The incoming filename should " "have the name .in, while the " "outgoing name should be .out")); +static cl::opt InteractiveIncludeDefault( + "inliner-interactive-include-default", cl::Hidden, + cl::desc( + (Twine("In interactive mode, also send the default policy decision: ") + + DefaultDecisionName + ".") + .str())); #if defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL) // codegen-ed file @@ -48,7 +54,8 @@ #endif std::unique_ptr -llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) { +llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM, + std::function GetDefaultAdvice) { if (!llvm::isEmbeddedModelEvaluatorValid() && InteractiveChannelBaseName.empty()) return nullptr; @@ -56,12 +63,17 @@ if (InteractiveChannelBaseName.empty()) AOTRunner = std::make_unique>( M.getContext(), FeatureMap, DecisionName); - else + else { + auto Features = FeatureMap; + if (InteractiveIncludeDefault) + Features.push_back(DefaultDecisionSpec); AOTRunner = std::make_unique( - M.getContext(), FeatureMap, InlineDecisionSpec, + M.getContext(), Features, InlineDecisionSpec, InteractiveChannelBaseName + ".out", InteractiveChannelBaseName + ".in"); - return std::make_unique(M, MAM, std::move(AOTRunner)); + } + return std::make_unique(M, MAM, std::move(AOTRunner), + GetDefaultAdvice); } #define DEBUG_TYPE "inline-ml" @@ -96,6 +108,8 @@ const TensorSpec llvm::InlineDecisionSpec = TensorSpec::createSpec(DecisionName, {1}); const char *const llvm::DefaultDecisionName = "inlining_default"; +const TensorSpec llvm::DefaultDecisionSpec = + TensorSpec::createSpec(DefaultDecisionName, {1}); const char *const llvm::RewardName = "delta_size"; CallBase *getInlinableCS(Instruction &I) { @@ -108,11 +122,13 @@ return nullptr; } -MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr Runner) +MLInlineAdvisor::MLInlineAdvisor( + Module &M, ModuleAnalysisManager &MAM, + std::unique_ptr Runner, + std::function GetDefaultAdvice) : InlineAdvisor( M, MAM.getResult(M).getManager()), - ModelRunner(std::move(Runner)), + ModelRunner(std::move(Runner)), GetDefaultAdvice(GetDefaultAdvice), CG(MAM.getResult(M)), InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize) { assert(ModelRunner); @@ -393,7 +409,10 @@ *ModelRunner->getTensor(inlineCostFeatureToMlFeature( static_cast(I))) = CostFeatures->at(I); } - + // This one would have been set up to be right at the end. + if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault) + *ModelRunner->getTensor(InlineCostFeatureIndex::NumberOfFeatures) = + GetDefaultAdvice(CB); return getAdviceFromModel(CB, ORE); } diff --git a/llvm/test/Transforms/Inline/ML/interactive-mode.ll b/llvm/test/Transforms/Inline/ML/interactive-mode.ll --- a/llvm/test/Transforms/Inline/ML/interactive-mode.ll +++ b/llvm/test/Transforms/Inline/ML/interactive-mode.ll @@ -7,7 +7,12 @@ ; RUN: cp %S/Inputs/interactive_main.py %t.rundir ; RUN: %python %t.rundir/interactive_main.py %t.channel-basename \ ; RUN: opt -passes=scc-oz-module-inliner -interactive-model-runner-echo-reply \ -; RUN: -enable-ml-inliner=release --inliner-interactive-channel-base=%t.channel-basename %S/Inputs/test-module.ll -S -o /dev/null | FileCheck %s +; RUN: -enable-ml-inliner=release -inliner-interactive-channel-base=%t.channel-basename %S/Inputs/test-module.ll -S -o /dev/null | FileCheck %s +; RUN: %python %t.rundir/interactive_main.py %t.channel-basename \ +; RUN: opt -passes=scc-oz-module-inliner -interactive-model-runner-echo-reply \ +; RUN: -inliner-interactive-include-default -enable-ml-inliner=release \ +; RUN: -inliner-interactive-channel-base=%t.channel-basename %S/Inputs/test-module.ll -S -o /dev/null | FileCheck %s -check-prefixes=CHECK,CHECK-DEFAULT + ;; It'd be nice if we had stdout and stderr interleaved, but we don't, so ;; let's just check the features have non-zero values, and that we see as many @@ -17,6 +22,7 @@ ; CHECK-NEXT: sroa_savings: 0 ; CHECK: unsimplified_common_instructions: 5 ; CHECK: callee_users: 3 +; CHECK-DEFAULT: inlining_default: 0 ; CHECK: observation: 5 ; CHECK-NOT: observation: 6