[mlgo][inl] Interactive mode: optionally tell the default decision

[mtrofin]

Summary

This helps training algorithms that may want to sometimes replicate the
default decision. The default decision is presented as an extra feature
called inlining_default. It's not normally exported to save
computation time.

This is only available in interactive mode.

[jacobhegna]

LGTM

[mehdi_amini]

llvm/lib/Analysis/MLInlineAdvisor.cpp

46

This leads to a use-after-free: you're creating a temporary string that is immediately destroyed but referenced through a StringRef when printing the help.

I pushed a quick fix in 437dfa5b0365 but you may want to tweak this.

[mtrofin]

Oh - duh. Thank you!

llvm/include/llvm/Analysis/InlineAdvisor.h

public:
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 
   PreservedAnalyses run(LazyCallGraph::SCC &InitialC, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR);
 };
 
 std::unique_ptr<InlineAdvisor>
-getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM);
+getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
+                      std::function<bool(CallBase &)> GetDefaultAdvice);
 
 std::unique_ptr<InlineAdvisor>
 getDevelopmentModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
                           std::function<bool(CallBase &)> GetDefaultAdvice);
 
 // Default (manual policy) decision making helper APIs. Shared with the legacy
 // pass manager inliner.
 

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

 constexpr size_t NumberOfFeatures =
     static_cast<size_t>(FeatureIndex::NumberOfFeatures);
 
 extern const std::vector<TensorSpec> FeatureMap;
 
 extern const char *const DecisionName;
 extern const TensorSpec InlineDecisionSpec;
 extern const char *const DefaultDecisionName;
+extern const TensorSpec DefaultDecisionSpec;
 extern const char *const RewardName;
 
 using InlineFeatures = std::vector<int64_t>;
 
 } // namespace llvm
 #endif // LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H

llvm/include/llvm/Analysis/MLInlineAdvisor.h

 namespace llvm {
 class DiagnosticInfoOptimizationBase;
 class Module;
 class MLInlineAdvice;
 
 class MLInlineAdvisor : public InlineAdvisor {
 public:
   MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
-                  std::unique_ptr<MLModelRunner> ModelRunner);
+                  std::unique_ptr<MLModelRunner> ModelRunner,
+                  std::function<bool(CallBase &)> GetDefaultAdvice);
 
   virtual ~MLInlineAdvisor() = default;
 
   void onPassEntry(LazyCallGraph::SCC *SCC) override;
   void onPassExit(LazyCallGraph::SCC *SCC) override;
 
   int64_t getIRSize(Function &F) const {
     return getCachedFPI(F).TotalInstructionCount;
protected:
   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE);
 
   // Get the initial 'level' of the function, or 0 if the function has been
   // introduced afterwards.
   // TODO: should we keep this updated?
   unsigned getInitialFunctionLevel(const Function &F) const;
 
   std::unique_ptr<MLModelRunner> ModelRunner;
+  std::function<bool(CallBase &)> GetDefaultAdvice;
 
 private:
   int64_t getModuleIRSize() const;
   std::unique_ptr<InlineAdvice>
   getSkipAdviceIfUnreachableCallsite(CallBase &CB);
   void print(raw_ostream &OS) const override;
 
   // Using std::map to benefit from its iterator / reference non-invalidating

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

public:
   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
 
   std::optional<size_t> getNativeSizeEstimate(const Function &F) const;
 
 private:
   bool isLogging() const { return !!Logger; }
   std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB) override;
 
-  std::function<bool(CallBase &)> GetDefaultAdvice;
   const bool IsDoingInference;
   std::unique_ptr<TrainingLogger> Logger;
 
   const std::optional<int32_t> InitialNativeSize;
   std::optional<int32_t> CurrentNativeSize;
 };
 
 /// A variant of MLInlineAdvice that tracks all non-trivial inlining
TrainingLogger::TrainingLogger(StringRef LogFileName,
     : LogFileName(LogFileName), MUTR(MUTR) {
   // The first output is the inlining decision.
   std::vector<TensorSpec> FT(FeatureMap.begin(), FeatureMap.end());
 
   if (MUTR)
     append_range(FT, MUTR->extraOutputsForLoggingSpecs());
 
   DefaultDecisionPos = FT.size();
-  FT.push_back(TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}));
+  FT.push_back(DefaultDecisionSpec);
 
   DecisionPos = FT.size();
   FT.push_back(InlineDecisionSpec);
   std::error_code EC;
   auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
   if (EC)
     dbgs() << (EC.message() + ":" + TrainingLog);
 
void TrainingLogger::logInlineEvent(const InlineEvent &Event,
   Effects.push_back(Event.Effect);
 }
 
 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(GetDefaultAdvice),
+    : MLInlineAdvisor(M, MAM, std::move(ModelRunner), 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.
   assert(IsDoingInference || isLogging());
 }
 

llvm/lib/Analysis/InlineAdvisor.cpp

bool InlineAdvisorAnalysis::Result::tryCreate(
     InlineParams Params, InliningAdvisorMode Mode,
     const ReplayInlinerSettings &ReplaySettings, InlineContext IC) {
   auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   if (PluginInlineAdvisorAnalysis::HasBeenRegistered) {
     auto &DA = MAM.getResult<PluginInlineAdvisorAnalysis>(M);
     Advisor.reset(DA.Factory(M, FAM, Params, IC));
     return !!Advisor;
   }
+  auto GetDefaultAdvice = [&FAM, Params](CallBase &CB) {
+    auto OIC = getDefaultInlineAdvice(CB, FAM, Params);
+    return OIC.has_value();
+  };
   switch (Mode) {
   case InliningAdvisorMode::Default:
     LLVM_DEBUG(dbgs() << "Using default inliner heuristic.\n");
     Advisor.reset(new DefaultInlineAdvisor(M, FAM, Params, IC));
     // Restrict replay to default advisor, ML advisors are stateful so
     // replay will need augmentations to interleave with them correctly.
     if (!ReplaySettings.ReplayFile.empty()) {
       Advisor = llvm::getReplayInlineAdvisor(M, FAM, M.getContext(),
                                              std::move(Advisor), ReplaySettings,
                                              /* EmitRemarks =*/true, IC);
     }
     break;
   case InliningAdvisorMode::Development:
 #ifdef LLVM_HAVE_TFLITE
     LLVM_DEBUG(dbgs() << "Using development-mode inliner policy.\n");
-    Advisor =
-        llvm::getDevelopmentModeAdvisor(M, MAM, [&FAM, Params](CallBase &CB) {
-          auto OIC = getDefaultInlineAdvice(CB, FAM, Params);
-          return OIC.has_value();
-        });
+    Advisor = llvm::getDevelopmentModeAdvisor(M, MAM, GetDefaultAdvice);
 #endif
     break;
   case InliningAdvisorMode::Release:
     LLVM_DEBUG(dbgs() << "Using release-mode inliner policy.\n");
-    Advisor = llvm::getReleaseModeAdvisor(M, MAM);
+    Advisor = llvm::getReleaseModeAdvisor(M, MAM, GetDefaultAdvice);
     break;
   }
 
   return !!Advisor;
 }
 
 /// Return true if inlining of CB can block the caller from being
 /// inlined which is proved to be more beneficial. \p IC is the

llvm/lib/Analysis/MLInlineAdvisor.cpp

 using namespace llvm;
 
 static cl::opt<std::string> InteractiveChannelBaseName(
     "inliner-interactive-channel-base", cl::Hidden,
     cl::desc(
         "Base file path for the interactive mode. The incoming filename should "
         "have the name <inliner-interactive-channel-base>.in, while the "
         "outgoing name should be <inliner-interactive-channel-base>.out"));
+static cl::opt<bool> InteractiveIncludeDefault(
+    "inliner-interactive-include-default", cl::Hidden,
+    cl::desc(
+        (Twine("In interactive mode, also send the default policy decision: ") +
+         DefaultDecisionName + ".")
+            .str()));
 
 #if defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
 // codegen-ed file
 #include "InlinerSizeModel.h" // NOLINT
 using CompiledModelType = llvm::InlinerSizeModel;
 #else
 using CompiledModelType = NoopSavedModelImpl;
 #endif
 
 std::unique_ptr<InlineAdvisor>
-llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
+llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
+                            std::function<bool(CallBase &)> GetDefaultAdvice) {
   if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
       InteractiveChannelBaseName.empty())
     return nullptr;
   std::unique_ptr<MLModelRunner> AOTRunner;
   if (InteractiveChannelBaseName.empty())
     AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
         M.getContext(), FeatureMap, DecisionName);
-  else
+  else {
+    auto Features = FeatureMap;
+    if (InteractiveIncludeDefault)
+      Features.push_back(DefaultDecisionSpec);
     AOTRunner = std::make_unique<InteractiveModelRunner>(
-        M.getContext(), FeatureMap, InlineDecisionSpec,
+        M.getContext(), Features, InlineDecisionSpec,
         InteractiveChannelBaseName + ".out",
         InteractiveChannelBaseName + ".in");
-  return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
+  }
+  return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner),
+                                           GetDefaultAdvice);
 }
 
 #define DEBUG_TYPE "inline-ml"
 
 static cl::opt<float> SizeIncreaseThreshold(
     "ml-advisor-size-increase-threshold", cl::Hidden,
     cl::desc("Maximum factor by which expected native size may increase before "
              "blocking any further inlining."),

 #undef POPULATE_NAMES
 };
 // clang-format on
 
 const char *const llvm::DecisionName = "inlining_decision";
 const TensorSpec llvm::InlineDecisionSpec =
     TensorSpec::createSpec<int64_t>(DecisionName, {1});
 const char *const llvm::DefaultDecisionName = "inlining_default";
+const TensorSpec llvm::DefaultDecisionSpec =
+    TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1});
 const char *const llvm::RewardName = "delta_size";
 
 CallBase *getInlinableCS(Instruction &I) {
   if (auto *CS = dyn_cast<CallBase>(&I))
     if (Function *Callee = CS->getCalledFunction()) {
       if (!Callee->isDeclaration()) {
         return CS;
       }
     }
   return nullptr;
 }
 
-MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
-                                 std::unique_ptr<MLModelRunner> Runner)
+MLInlineAdvisor::MLInlineAdvisor(
+    Module &M, ModuleAnalysisManager &MAM,
+    std::unique_ptr<MLModelRunner> Runner,
+    std::function<bool(CallBase &)> GetDefaultAdvice)
     : InlineAdvisor(
           M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
-      ModelRunner(std::move(Runner)),
+      ModelRunner(std::move(Runner)), GetDefaultAdvice(GetDefaultAdvice),
       CG(MAM.getResult<LazyCallGraphAnalysis>(M)),
       InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize) {
   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
   // critical in behavioral cloning - i.e. training a model to mimic the manual
std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   *ModelRunner->getTensor<int64_t>(FeatureIndex::CostEstimate) = CostEstimate;
 
   // Add the cost features
   for (size_t I = 0;
        I < static_cast<size_t>(InlineCostFeatureIndex::NumberOfFeatures); ++I) {
     *ModelRunner->getTensor<int64_t>(inlineCostFeatureToMlFeature(
         static_cast<InlineCostFeatureIndex>(I))) = CostFeatures->at(I);
   }
-
+  // This one would have been set up to be right at the end.
+  if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault)
+    *ModelRunner->getTensor<int64_t>(InlineCostFeatureIndex::NumberOfFeatures) =
+        GetDefaultAdvice(CB);
   return getAdviceFromModel(CB, ORE);
 }
 
 std::unique_ptr<MLInlineAdvice>
 MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
                                     OptimizationRemarkEmitter &ORE) {
   return std::make_unique<MLInlineAdvice>(
       this, CB, ORE, static_cast<bool>(ModelRunner->evaluate<int64_t>()));

llvm/test/Transforms/Inline/ML/interactive-mode.ll

 ; REQUIRES: x86_64-linux
 ; RUN: rm -rf %t.rundir
 ; RUN: rm -rf %t.channel-basename.*
 ; RUN: mkdir %t.rundir
 ; RUN: cp %S/../../../../lib/Analysis/models/log_reader.py %t.rundir
 ; RUN: cp %S/../../../../lib/Analysis/models/interactive_host.py %t.rundir
 ; RUN: cp %S/Inputs/interactive_main.py %t.rundir
 ; RUN: %python %t.rundir/interactive_main.py %t.channel-basename \
 ; RUN:    opt -passes=scc-oz-module-inliner -interactive-model-runner-echo-reply \
-; RUN:    -enable-ml-inliner=release --inliner-interactive-channel-base=%t.channel-basename %S/Inputs/test-module.ll -S -o /dev/null | FileCheck %s
+; RUN:    -enable-ml-inliner=release -inliner-interactive-channel-base=%t.channel-basename %S/Inputs/test-module.ll -S -o /dev/null | FileCheck %s
+; RUN: %python %t.rundir/interactive_main.py %t.channel-basename \
+; RUN:    opt -passes=scc-oz-module-inliner -interactive-model-runner-echo-reply \
+; RUN:    -inliner-interactive-include-default -enable-ml-inliner=release \
+; RUN:    -inliner-interactive-channel-base=%t.channel-basename %S/Inputs/test-module.ll -S -o /dev/null | FileCheck %s -check-prefixes=CHECK,CHECK-DEFAULT
+
 
 ;; It'd be nice if we had stdout and stderr interleaved, but we don't, so
 ;; let's just check the features have non-zero values, and that we see as many
 ;; advices as observations, and that the advices flip-flop as intended.
 ; CHECK: context:
 ; CHECK-NEXT: observation: 0
 ; CHECK-NEXT: sroa_savings: 0
 ; CHECK:      unsimplified_common_instructions: 5
 ; CHECK:      callee_users: 3
+; CHECK-DEFAULT: inlining_default: 0
 ; CHECK:      observation: 5
 ; CHECK-NOT:  observation: 6
 
 ; CHECK:      inlining_decision: 1
 ; CHECK-NEXT: inlining_decision: 0
 ; CHECK-NEXT: inlining_decision: 1
 ; CHECK-NEXT: inlining_decision: 0
 ; CHECK-NEXT: inlining_decision: 1
 ; CHECK-NEXT: inlining_decision: 0