[mlgo][inliner] Fix potential concurrency issue in local ThinLTO + IR2Vec cases (#156120)

mtrofin · web-flow · commit c878baf1d9a2 · 2025-08-29T18:24:30.000-07:00
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)
diff --git 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
@@ -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
@@ -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
@@ -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());
 }

Original file line number	Diff line number	Diff line change
`@@ -160,8 +160,6 @@ inlineCostFeatureToMlFeature(InlineCostFeatureIndex Feature) {`
`160`	`160`	`return static_cast<FeatureIndex>(static_cast<size_t>(Feature));`
`161`	`161`	`}`
`162`	`162`
`163`		`-LLVM_ABI extern std::vector<TensorSpec> &getFeatureMap();`
`164`		`-`
`165`	`163`	`LLVM_ABI extern const char *const DecisionName;`
`166`	`164`	`LLVM_ABI extern const TensorSpec InlineDecisionSpec;`
`167`	`165`	`LLVM_ABI extern const char *const DefaultDecisionName;`