[mlgo][inliner] Fix potential concurrency issue in local ThinLTO + IR2Vec cases #156120
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
…2Vec cases The inliner's `FeatureMap` used to be immutable, but in IR2Vec cases we don't know the shapes of the embedding vectors until later, so we need to initialize it at the time we construct the advisor. In non-distributed ThinLTO cases, for example, this means we'd mutate shared state. The feature set is also needed when constructing the underlying model runner. The alternative here is to postpone the creation of the model runner to the time we construct the advisor, and also make the feature map a member of the advisor object.
llvmbot added mlgo llvm:analysis 
mtrofin changed the title 
Member
| @llvm/pr-subscribers-mlgo Author: Mircea Trofin (mtrofin) ChangesThe inliner's The feature set is also needed when constructing the underlying model runner. The alternative here is to postpone the creation of the model runner to the time we construct the advisor, and also make the feature map a member of the advisor object. Full diff: https://github.com/llvm/llvm-project/pull/156120.diff 4 Files Affected:
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<FeatureIndex>(static_cast<size_t>(Feature)); } -LLVM_ABI extern std::vector<TensorSpec> &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<MLModelRunner> ModelRunner, + std::function<std::unique_ptr<MLModelRunner>( + const std::vector<TensorSpec> &)> + GetModelRunner, std::function<bool(CallBase &)> 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<TensorSpec> &getFeatureMap() const { return FeatureMap; }; + static const std::vector<TensorSpec> &getInitialFeatureMap(); protected: std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override; @@ -65,6 +69,7 @@ class MLInlineAdvisor : public InlineAdvisor { std::unique_ptr<MLModelRunner> ModelRunner; std::function<bool(CallBase &)> GetDefaultAdvice; + std::vector<TensorSpec> FeatureMap; private: int64_t getModuleIRSize() const; diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp index 790e00e1b3b06..99cd7364a4618 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<TensorSpec> &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<TensorSpec> &FeatureMap; + std::unique_ptr<Logger> 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<MLModelRunner> ModelRunner, - std::function<bool(CallBase &)> GetDefaultAdvice, - std::unique_ptr<TrainingLogger> Logger); + std::function< + std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)> + GetModelRunner, + std::function<bool(CallBase &)> GetDefaultAdvice); size_t getTotalSizeEstimate(); @@ -258,9 +262,13 @@ static const std::vector<TensorSpec> TrainingOnlyFeatures{ TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}), TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})}; -static const std::vector<TensorSpec> 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<TensorSpec> +convertInputFeatures(const std::vector<TensorSpec> &OriginalFeatures) { std::vector<TensorSpec> InputSpecs; - for (const auto &Feature : FeatureMap) + 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<TensorSpec> getInputFeatures() { } // namespace TrainingLogger::TrainingLogger(StringRef LogFileName, - const ModelUnderTrainingRunner *MUTR) - : LogFileName(LogFileName), MUTR(MUTR) { + const ModelUnderTrainingRunner *MUTR, + const std::vector<TensorSpec> &FeatureMap) + : LogFileName(LogFileName), MUTR(MUTR), FeatureMap(FeatureMap) { // The first output is the inlining decision. std::vector<TensorSpec> FT(FeatureMap.begin(), FeatureMap.end()); @@ -327,15 +336,19 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event, DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr<MLModelRunner> ModelRunner, - std::function<bool(CallBase &)> GetDefaultAdvice, - std::unique_ptr<TrainingLogger> Logger) - : MLInlineAdvisor(M, MAM, std::move(ModelRunner), GetDefaultAdvice), + std::function< + std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)> + GetModelRunner, + std::function<bool(CallBase &)> GetDefaultAdvice) + : MLInlineAdvisor(M, MAM, GetModelRunner, GetDefaultAdvice), IsDoingInference(isa<ModelUnderTrainingRunner>(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<TrainingLogger>( + TrainingLog, dyn_cast<ModelUnderTrainingRunner>(ModelRunner.get()), + getFeatureMap()); assert(IsDoingInference || isLogging()); } @@ -401,21 +414,22 @@ std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor( Module &M, ModuleAnalysisManager &MAM, std::function<bool(CallBase &)> GetDefaultAdvice) { auto &Ctx = M.getContext(); - std::unique_ptr<MLModelRunner> 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<TrainingLogger> Logger; - if (!TrainingLog.empty()) - Logger = std::make_unique<TrainingLogger>( - TrainingLog, dyn_cast<ModelUnderTrainingRunner>(Runner.get())); - - return std::make_unique<DevelopmentModeMLInlineAdvisor>( - M, MAM, std::move(Runner), GetDefaultAdvice, std::move(Logger)); + auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures) + -> std::unique_ptr<MLModelRunner> { + std::unique_ptr<MLModelRunner> Runner; + const std::vector<TensorSpec> 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<DevelopmentModeMLInlineAdvisor>(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<CompiledModelType>() && InteractiveChannelBaseName.empty()) return nullptr; - std::unique_ptr<MLModelRunner> AOTRunner; - if (InteractiveChannelBaseName.empty()) - AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>( - M.getContext(), getFeatureMap(), DecisionName, - EmbeddedModelRunnerOptions().setModelSelector(ModelSelector)); - else { - auto Features = getFeatureMap(); - if (InteractiveIncludeDefault) - Features.push_back(DefaultDecisionSpec); - AOTRunner = std::make_unique<InteractiveModelRunner>( - M.getContext(), Features, InlineDecisionSpec, - InteractiveChannelBaseName + ".out", - InteractiveChannelBaseName + ".in"); - } - return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner), + auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures) + -> std::unique_ptr<MLModelRunner> { + std::unique_ptr<MLModelRunner> AOTRunner; + if (InteractiveChannelBaseName.empty()) + AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>( + M.getContext(), InputFeatures, DecisionName, + EmbeddedModelRunnerOptions().setModelSelector(ModelSelector)); + else { + AOTRunner = std::make_unique<InteractiveModelRunner>( + M.getContext(), InputFeatures, InlineDecisionSpec, + InteractiveChannelBaseName + ".out", + InteractiveChannelBaseName + ".in"); + } + return AOTRunner; + }; + return std::make_unique<MLInlineAdvisor>(M, MAM, RunnerFactory, GetDefaultAdvice); } @@ -107,7 +108,7 @@ static cl::opt<bool> KeepFPICache( "For test - keep the ML Inline advisor's FunctionPropertiesInfo cache"), cl::init(false)); -std::vector<TensorSpec> &llvm::getFeatureMap() { +const std::vector<TensorSpec> &MLInlineAdvisor::getInitialFeatureMap() { // clang-format off static std::vector<TensorSpec> FeatureMap{ #define POPULATE_NAMES(DTYPE, SHAPE, NAME, __) TensorSpec::createSpec<DTYPE>(#NAME, SHAPE), @@ -142,17 +143,17 @@ CallBase *getInlinableCS(Instruction &I) { MLInlineAdvisor::MLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr<MLModelRunner> Runner, + std::function< + std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)> + GetModelRunner, std::function<bool(CallBase &)> GetDefaultAdvice) : InlineAdvisor( M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()), - ModelRunner(std::move(Runner)), GetDefaultAdvice(GetDefaultAdvice), + GetDefaultAdvice(GetDefaultAdvice), FeatureMap(getInitialFeatureMap()), CG(MAM.getResult<LazyCallGraphAnalysis>(M)), UseIR2Vec(MAM.getCachedResult<IR2VecVocabAnalysis>(M) != nullptr), InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize), PSI(MAM.getResult<ProfileSummaryAnalysis>(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<IR2VecVocabAnalysis>(M)) { + if (auto *IR2VecVocabResult = MAM.getCachedResult<IR2VecVocabAnalysis>(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<float>("callee_embedding", {IR2VecDim})); - getFeatureMap().push_back( + FeatureMap.push_back( TensorSpec::createSpec<float>("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<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) { } // This one would have been set up to be right at the end. if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault) - *ModelRunner->getTensor<int64_t>(getFeatureMap().size()) = + *ModelRunner->getTensor<int64_t>(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<int64_t>(I)); OR << NV("ShouldInline", isInliningRecommended()); } |
Member
| @llvm/pr-subscribers-llvm-analysis Author: Mircea Trofin (mtrofin) ChangesThe inliner's The feature set is also needed when constructing the underlying model runner. The alternative here is to postpone the creation of the model runner to the time we construct the advisor, and also make the feature map a member of the advisor object. Full diff: https://github.com/llvm/llvm-project/pull/156120.diff 4 Files Affected:
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<FeatureIndex>(static_cast<size_t>(Feature)); } -LLVM_ABI extern std::vector<TensorSpec> &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<MLModelRunner> ModelRunner, + std::function<std::unique_ptr<MLModelRunner>( + const std::vector<TensorSpec> &)> + GetModelRunner, std::function<bool(CallBase &)> 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<TensorSpec> &getFeatureMap() const { return FeatureMap; }; + static const std::vector<TensorSpec> &getInitialFeatureMap(); protected: std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override; @@ -65,6 +69,7 @@ class MLInlineAdvisor : public InlineAdvisor { std::unique_ptr<MLModelRunner> ModelRunner; std::function<bool(CallBase &)> GetDefaultAdvice; + std::vector<TensorSpec> FeatureMap; private: int64_t getModuleIRSize() const; diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp index 790e00e1b3b06..99cd7364a4618 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<TensorSpec> &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<TensorSpec> &FeatureMap; + std::unique_ptr<Logger> 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<MLModelRunner> ModelRunner, - std::function<bool(CallBase &)> GetDefaultAdvice, - std::unique_ptr<TrainingLogger> Logger); + std::function< + std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)> + GetModelRunner, + std::function<bool(CallBase &)> GetDefaultAdvice); size_t getTotalSizeEstimate(); @@ -258,9 +262,13 @@ static const std::vector<TensorSpec> TrainingOnlyFeatures{ TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}), TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})}; -static const std::vector<TensorSpec> 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<TensorSpec> +convertInputFeatures(const std::vector<TensorSpec> &OriginalFeatures) { std::vector<TensorSpec> InputSpecs; - for (const auto &Feature : FeatureMap) + 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<TensorSpec> getInputFeatures() { } // namespace TrainingLogger::TrainingLogger(StringRef LogFileName, - const ModelUnderTrainingRunner *MUTR) - : LogFileName(LogFileName), MUTR(MUTR) { + const ModelUnderTrainingRunner *MUTR, + const std::vector<TensorSpec> &FeatureMap) + : LogFileName(LogFileName), MUTR(MUTR), FeatureMap(FeatureMap) { // The first output is the inlining decision. std::vector<TensorSpec> FT(FeatureMap.begin(), FeatureMap.end()); @@ -327,15 +336,19 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event, DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr<MLModelRunner> ModelRunner, - std::function<bool(CallBase &)> GetDefaultAdvice, - std::unique_ptr<TrainingLogger> Logger) - : MLInlineAdvisor(M, MAM, std::move(ModelRunner), GetDefaultAdvice), + std::function< + std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)> + GetModelRunner, + std::function<bool(CallBase &)> GetDefaultAdvice) + : MLInlineAdvisor(M, MAM, GetModelRunner, GetDefaultAdvice), IsDoingInference(isa<ModelUnderTrainingRunner>(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<TrainingLogger>( + TrainingLog, dyn_cast<ModelUnderTrainingRunner>(ModelRunner.get()), + getFeatureMap()); assert(IsDoingInference || isLogging()); } @@ -401,21 +414,22 @@ std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor( Module &M, ModuleAnalysisManager &MAM, std::function<bool(CallBase &)> GetDefaultAdvice) { auto &Ctx = M.getContext(); - std::unique_ptr<MLModelRunner> 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<TrainingLogger> Logger; - if (!TrainingLog.empty()) - Logger = std::make_unique<TrainingLogger>( - TrainingLog, dyn_cast<ModelUnderTrainingRunner>(Runner.get())); - - return std::make_unique<DevelopmentModeMLInlineAdvisor>( - M, MAM, std::move(Runner), GetDefaultAdvice, std::move(Logger)); + auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures) + -> std::unique_ptr<MLModelRunner> { + std::unique_ptr<MLModelRunner> Runner; + const std::vector<TensorSpec> 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<DevelopmentModeMLInlineAdvisor>(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<CompiledModelType>() && InteractiveChannelBaseName.empty()) return nullptr; - std::unique_ptr<MLModelRunner> AOTRunner; - if (InteractiveChannelBaseName.empty()) - AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>( - M.getContext(), getFeatureMap(), DecisionName, - EmbeddedModelRunnerOptions().setModelSelector(ModelSelector)); - else { - auto Features = getFeatureMap(); - if (InteractiveIncludeDefault) - Features.push_back(DefaultDecisionSpec); - AOTRunner = std::make_unique<InteractiveModelRunner>( - M.getContext(), Features, InlineDecisionSpec, - InteractiveChannelBaseName + ".out", - InteractiveChannelBaseName + ".in"); - } - return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner), + auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures) + -> std::unique_ptr<MLModelRunner> { + std::unique_ptr<MLModelRunner> AOTRunner; + if (InteractiveChannelBaseName.empty()) + AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>( + M.getContext(), InputFeatures, DecisionName, + EmbeddedModelRunnerOptions().setModelSelector(ModelSelector)); + else { + AOTRunner = std::make_unique<InteractiveModelRunner>( + M.getContext(), InputFeatures, InlineDecisionSpec, + InteractiveChannelBaseName + ".out", + InteractiveChannelBaseName + ".in"); + } + return AOTRunner; + }; + return std::make_unique<MLInlineAdvisor>(M, MAM, RunnerFactory, GetDefaultAdvice); } @@ -107,7 +108,7 @@ static cl::opt<bool> KeepFPICache( "For test - keep the ML Inline advisor's FunctionPropertiesInfo cache"), cl::init(false)); -std::vector<TensorSpec> &llvm::getFeatureMap() { +const std::vector<TensorSpec> &MLInlineAdvisor::getInitialFeatureMap() { // clang-format off static std::vector<TensorSpec> FeatureMap{ #define POPULATE_NAMES(DTYPE, SHAPE, NAME, __) TensorSpec::createSpec<DTYPE>(#NAME, SHAPE), @@ -142,17 +143,17 @@ CallBase *getInlinableCS(Instruction &I) { MLInlineAdvisor::MLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, - std::unique_ptr<MLModelRunner> Runner, + std::function< + std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)> + GetModelRunner, std::function<bool(CallBase &)> GetDefaultAdvice) : InlineAdvisor( M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()), - ModelRunner(std::move(Runner)), GetDefaultAdvice(GetDefaultAdvice), + GetDefaultAdvice(GetDefaultAdvice), FeatureMap(getInitialFeatureMap()), CG(MAM.getResult<LazyCallGraphAnalysis>(M)), UseIR2Vec(MAM.getCachedResult<IR2VecVocabAnalysis>(M) != nullptr), InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize), PSI(MAM.getResult<ProfileSummaryAnalysis>(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<IR2VecVocabAnalysis>(M)) { + if (auto *IR2VecVocabResult = MAM.getCachedResult<IR2VecVocabAnalysis>(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<float>("callee_embedding", {IR2VecDim})); - getFeatureMap().push_back( + FeatureMap.push_back( TensorSpec::createSpec<float>("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<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) { } // This one would have been set up to be right at the end. if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault) - *ModelRunner->getTensor<int64_t>(getFeatureMap().size()) = + *ModelRunner->getTensor<int64_t>(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<int64_t>(I)); OR << NV("ShouldInline", isInliningRecommended()); } |
svkeerthy reviewed Aug 29, 2025
| FeatureMap.push_back( | ||
| TensorSpec::createSpec<float>("caller_embedding", {IR2VecDim})); | ||
| } | ||
| if (InteractiveIncludeDefault) |
Contributor
There was a problem hiding this comment.
Should it also check for !InteractiveChannelBaseName.empty()?
Member Author
There was a problem hiding this comment.
ya, but if not, and the model doesn't have it, it's a noop (minus compiletime penalty).
svkeerthy reviewed Aug 29, 2025
| // This one would have been set up to be right at the end. | ||
| if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault) | ||
| *ModelRunner->getTensor<int64_t>(getFeatureMap().size()) = | ||
| *ModelRunner->getTensor<int64_t>(getFeatureMap().size() - 1) = |
Contributor
There was a problem hiding this comment.
Trying to understand.. why did this change?
Member Author
There was a problem hiding this comment.
Before, we were adding the extra feature to a local clone we'd send the model (see lines 84-85 on the left)
boomanaiden154 approved these changes Aug 30, 2025
Collaborator
| LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/137/builds/24546 Here is the relevant piece of the build log for the reference |
Collaborator
| LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/175/builds/24385 Here is the relevant piece of the build log for the reference |
mtrofin added a commit to mtrofin/llvm-project that referenced this pull request Aug 30, 2025
mtrofin added a commit that referenced this pull request Aug 30, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.
The inliner's
FeatureMapused to be immutable, but in IR2Vec cases we don't know the shapes of the embedding vectors until later, so we need to initialize it at the time we construct the advisor. In non-distributed ThinLTO cases, for example, this means we'd mutate shared state.The feature set is also needed when constructing the underlying model runner.
The alternative here is to postpone the creation of the model runner to the time we construct the advisor, and also make the feature map a member of the advisor object.
(issue identified by @efriedma-quic in PR #154541)