diff --git a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h index 5c6aee3ab38ab..e559171b9c257 100644 --- a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h +++ b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h @@ -160,8 +160,6 @@ inlineCostFeatureToMlFeature(InlineCostFeatureIndex Feature) { return static_cast(static_cast(Feature)); } -LLVM_ABI extern std::vector &getFeatureMap(); - LLVM_ABI extern const char *const DecisionName; LLVM_ABI extern const TensorSpec InlineDecisionSpec; LLVM_ABI extern const char *const DefaultDecisionName; diff --git a/llvm/include/llvm/Analysis/MLInlineAdvisor.h b/llvm/include/llvm/Analysis/MLInlineAdvisor.h index 8262dd0846ede..cc4c482b379e3 100644 --- a/llvm/include/llvm/Analysis/MLInlineAdvisor.h +++ b/llvm/include/llvm/Analysis/MLInlineAdvisor.h @@ -28,7 +28,9 @@ class ProfileSummaryInfo; class MLInlineAdvisor : public InlineAdvisor { public: MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr ModelRunner, + std::function( + const std::vector &)> + GetModelRunner, std::function GetDefaultAdvice); virtual ~MLInlineAdvisor() = default; @@ -46,6 +48,8 @@ class MLInlineAdvisor : public InlineAdvisor { int64_t getLocalCalls(Function &F); const MLModelRunner &getModelRunner() const { return *ModelRunner; } FunctionPropertiesInfo &getCachedFPI(Function &) const; + const std::vector &getFeatureMap() const { return FeatureMap; }; + static const std::vector &getInitialFeatureMap(); protected: std::unique_ptr getAdviceImpl(CallBase &CB) override; @@ -65,6 +69,7 @@ class MLInlineAdvisor : public InlineAdvisor { std::unique_ptr ModelRunner; std::function GetDefaultAdvice; + std::vector FeatureMap; private: int64_t getModuleIRSize() const; diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp index ce2d8b654bf21..a51f62fe65776 100644 --- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp +++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp @@ -97,7 +97,8 @@ struct InlineEvent { /// Collect data we may use for training a model. class TrainingLogger final { public: - TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR); + TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR, + const std::vector &FeatureMap); /// Log one inlining event. void logInlineEvent(const InlineEvent &Event, @@ -106,6 +107,8 @@ class TrainingLogger final { private: StringRef LogFileName; const ModelUnderTrainingRunner *const MUTR; + const std::vector &FeatureMap; + std::unique_ptr L; BitVector Effects; /// Set these 2 clearly OOB, to make sure we set them later. @@ -142,9 +145,10 @@ class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor { public: DevelopmentModeMLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr ModelRunner, - std::function GetDefaultAdvice, - std::unique_ptr Logger); + std::function< + std::unique_ptr(const std::vector &)> + GetModelRunner, + std::function GetDefaultAdvice); size_t getTotalSizeEstimate(); @@ -258,9 +262,13 @@ static const std::vector TrainingOnlyFeatures{ TensorSpec::createSpec(TFFeedPrefix + "reward", {1}), TensorSpec::createSpec(TFFeedPrefix + "step_type", {1})}; -static const std::vector getInputFeatures() { +// add TFFeedPrefix to the names and also add the "TrainingOnlyFeatures" which +// the model runner needs to see present. We don't set them ourselves or +// interact with them. +static const std::vector +convertInputFeatures(const std::vector &OriginalFeatures) { std::vector InputSpecs; - for (const auto &Feature : getFeatureMap()) + for (const auto &Feature : OriginalFeatures) InputSpecs.push_back(TensorSpec(TFFeedPrefix + Feature.name(), Feature)); append_range(InputSpecs, TrainingOnlyFeatures); return InputSpecs; @@ -269,8 +277,9 @@ static const std::vector getInputFeatures() { } // namespace TrainingLogger::TrainingLogger(StringRef LogFileName, - const ModelUnderTrainingRunner *MUTR) - : LogFileName(LogFileName), MUTR(MUTR) { + const ModelUnderTrainingRunner *MUTR, + const std::vector &FeatureMap) + : LogFileName(LogFileName), MUTR(MUTR), FeatureMap(FeatureMap) { // The first output is the inlining decision. std::vector FT(getFeatureMap().begin(), getFeatureMap().end()); @@ -327,15 +336,19 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event, DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr ModelRunner, - std::function GetDefaultAdvice, - std::unique_ptr Logger) - : MLInlineAdvisor(M, MAM, std::move(ModelRunner), GetDefaultAdvice), + std::function< + std::unique_ptr(const std::vector &)> + GetModelRunner, + std::function GetDefaultAdvice) + : MLInlineAdvisor(M, MAM, GetModelRunner, GetDefaultAdvice), IsDoingInference(isa(getModelRunner())), - Logger(std::move(Logger)), InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0), CurrentNativeSize(InitialNativeSize) { // We cannot have the case of neither inference nor logging. + if (!TrainingLog.empty()) + Logger = std::make_unique( + TrainingLog, dyn_cast(ModelRunner.get()), + getFeatureMap()); assert(IsDoingInference || isLogging()); } @@ -401,21 +414,22 @@ std::unique_ptr llvm::getDevelopmentModeAdvisor( Module &M, ModuleAnalysisManager &MAM, std::function GetDefaultAdvice) { auto &Ctx = M.getContext(); - std::unique_ptr Runner; - if (TFModelUnderTrainingPath.empty()) - Runner.reset(new NoInferenceModelRunner(Ctx, getInputFeatures())); - else - Runner = ModelUnderTrainingRunner::createAndEnsureValid( - Ctx, TFModelUnderTrainingPath, DecisionName, getInputFeatures(), - TFOutputSpecOverride); - if (!Runner) - return nullptr; - std::unique_ptr Logger; - if (!TrainingLog.empty()) - Logger = std::make_unique( - TrainingLog, dyn_cast(Runner.get())); - - return std::make_unique( - M, MAM, std::move(Runner), GetDefaultAdvice, std::move(Logger)); + auto RunnerFactory = [&](const std::vector &InputFeatures) + -> std::unique_ptr { + std::unique_ptr Runner; + const std::vector ConvertedFeatures = + convertInputFeatures(InputFeatures); + if (TFModelUnderTrainingPath.empty()) + Runner.reset(new NoInferenceModelRunner(Ctx, ConvertedFeatures)); + else + Runner = ModelUnderTrainingRunner::createAndEnsureValid( + Ctx, TFModelUnderTrainingPath, DecisionName, ConvertedFeatures, + TFOutputSpecOverride); + if (!Runner) + return nullptr; + return Runner; + }; + return std::make_unique(M, MAM, RunnerFactory, + GetDefaultAdvice); } #endif // defined(LLVM_HAVE_TFLITE) diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp index 7854c19088ad3..f90717d3085eb 100644 --- a/llvm/lib/Analysis/MLInlineAdvisor.cpp +++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp @@ -75,21 +75,22 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM, if (!llvm::isEmbeddedModelEvaluatorValid() && InteractiveChannelBaseName.empty()) return nullptr; - std::unique_ptr AOTRunner; - if (InteractiveChannelBaseName.empty()) - AOTRunner = std::make_unique>( - M.getContext(), getFeatureMap(), DecisionName, - EmbeddedModelRunnerOptions().setModelSelector(ModelSelector)); - else { - auto Features = getFeatureMap(); - if (InteractiveIncludeDefault) - Features.push_back(DefaultDecisionSpec); - AOTRunner = std::make_unique( - M.getContext(), Features, InlineDecisionSpec, - InteractiveChannelBaseName + ".out", - InteractiveChannelBaseName + ".in"); - } - return std::make_unique(M, MAM, std::move(AOTRunner), + auto RunnerFactory = [&](const std::vector &InputFeatures) + -> std::unique_ptr { + std::unique_ptr AOTRunner; + if (InteractiveChannelBaseName.empty()) + AOTRunner = std::make_unique>( + M.getContext(), InputFeatures, DecisionName, + EmbeddedModelRunnerOptions().setModelSelector(ModelSelector)); + else { + AOTRunner = std::make_unique( + M.getContext(), InputFeatures, InlineDecisionSpec, + InteractiveChannelBaseName + ".out", + InteractiveChannelBaseName + ".in"); + } + return AOTRunner; + }; + return std::make_unique(M, MAM, RunnerFactory, GetDefaultAdvice); } @@ -107,7 +108,7 @@ static cl::opt KeepFPICache( "For test - keep the ML Inline advisor's FunctionPropertiesInfo cache"), cl::init(false)); -std::vector &llvm::getFeatureMap() { +const std::vector &MLInlineAdvisor::getInitialFeatureMap() { // clang-format off static std::vector FeatureMap{ #define POPULATE_NAMES(DTYPE, SHAPE, NAME, __) TensorSpec::createSpec(#NAME, SHAPE), @@ -142,17 +143,17 @@ CallBase *getInlinableCS(Instruction &I) { MLInlineAdvisor::MLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr Runner, + std::function< + std::unique_ptr(const std::vector &)> + GetModelRunner, std::function GetDefaultAdvice) : InlineAdvisor( M, MAM.getResult(M).getManager()), - ModelRunner(std::move(Runner)), GetDefaultAdvice(GetDefaultAdvice), + GetDefaultAdvice(GetDefaultAdvice), FeatureMap(getInitialFeatureMap()), CG(MAM.getResult(M)), UseIR2Vec(MAM.getCachedResult(M) != nullptr), InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize), PSI(MAM.getResult(M)) { - assert(ModelRunner); - ModelRunner->switchContext(""); // Extract the 'call site height' feature - the position of a call site // relative to the farthest statically reachable SCC node. We don't mutate // this value while inlining happens. Empirically, this feature proved @@ -192,18 +193,27 @@ MLInlineAdvisor::MLInlineAdvisor( } NodeCount = AllNodes.size(); - if (auto IR2VecVocabResult = MAM.getCachedResult(M)) { + if (auto *IR2VecVocabResult = MAM.getCachedResult(M)) { if (!IR2VecVocabResult->isValid()) { M.getContext().emitError("IR2VecVocabAnalysis is not valid"); return; } // Add the IR2Vec features to the feature map auto IR2VecDim = IR2VecVocabResult->getDimension(); - getFeatureMap().push_back( + FeatureMap.push_back( TensorSpec::createSpec("callee_embedding", {IR2VecDim})); - getFeatureMap().push_back( + FeatureMap.push_back( TensorSpec::createSpec("caller_embedding", {IR2VecDim})); } + if (InteractiveIncludeDefault) + FeatureMap.push_back(DefaultDecisionSpec); + + ModelRunner = GetModelRunner(getFeatureMap()); + if (!ModelRunner) { + M.getContext().emitError("Could not create model runner"); + return; + } + ModelRunner->switchContext(""); } unsigned MLInlineAdvisor::getInitialFunctionLevel(const Function &F) const { @@ -475,7 +485,7 @@ std::unique_ptr MLInlineAdvisor::getAdviceImpl(CallBase &CB) { } // This one would have been set up to be right at the end. if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault) - *ModelRunner->getTensor(getFeatureMap().size()) = + *ModelRunner->getTensor(getFeatureMap().size() - 1) = GetDefaultAdvice(CB); return getAdviceFromModel(CB, ORE); } @@ -554,8 +564,8 @@ void MLInlineAdvice::reportContextForRemark( DiagnosticInfoOptimizationBase &OR) { using namespace ore; OR << NV("Callee", Callee->getName()); - for (size_t I = 0; I < getFeatureMap().size(); ++I) - OR << NV(getFeatureMap()[I].name(), + for (size_t I = 0; I < getAdvisor()->getFeatureMap().size(); ++I) + OR << NV(getAdvisor()->getFeatureMap()[I].name(), *getAdvisor()->getModelRunner().getTensor(I)); OR << NV("ShouldInline", isInliningRecommended()); }