[llvm] [mlgo][inliner] Fix potential concurrency issue in local ThinLTO + IR… (PR #156120)
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Fri Aug 29 15:47:43 PDT 2025
https://github.com/mtrofin created https://github.com/llvm/llvm-project/pull/156120
…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.
>From 0a1e883d4286ccbd8f2e0f19a6a14248280871c6 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Fri, 29 Aug 2025 22:41:24 +0000
Subject: [PATCH] [mlgo][inliner] Fix potential concurrency issue in local
ThinLTO + IR2Vec 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.
---
.../llvm/Analysis/InlineModelFeatureMaps.h | 2 -
llvm/include/llvm/Analysis/MLInlineAdvisor.h | 7 +-
.../Analysis/DevelopmentModeInlineAdvisor.cpp | 72 +++++++++++--------
llvm/lib/Analysis/MLInlineAdvisor.cpp | 62 +++++++++-------
4 files changed, 85 insertions(+), 58 deletions(-)
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());
}
More information about the llvm-commits
mailing list