[llvm] c878baf - [mlgo][inliner] Fix potential concurrency issue in local ThinLTO + IR2Vec cases (#156120)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Aug 29 18:24:33 PDT 2025
Author: Mircea Trofin
Date: 2025-08-29T18:24:30-07:00
New Revision: c878baf1d9a259cf0788ffa1cc5c9d065adcb4c5
URL: https://github.com/llvm/llvm-project/commit/c878baf1d9a259cf0788ffa1cc5c9d065adcb4c5
DIFF: https://github.com/llvm/llvm-project/commit/c878baf1d9a259cf0788ffa1cc5c9d065adcb4c5.diff
LOG: [mlgo][inliner] Fix potential concurrency issue in local ThinLTO + IR2Vec cases (#156120)
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.
(issue identified by @efriedma-quic in PR #154541)
Added:
Modified:
llvm/include/llvm/Analysis/InlineModelFeatureMaps.h
llvm/include/llvm/Analysis/MLInlineAdvisor.h
llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
llvm/lib/Analysis/MLInlineAdvisor.cpp
Removed:
################################################################################
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 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<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 : 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<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(getFeatureMap().begin(), getFeatureMap().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());
}
More information about the llvm-commits
mailing list