[mlgo][inliner] Fix potential concurrency issue in local ThinLTO + IR2Vec cases

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;

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;

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)

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());
 }