[MLInliner] Support training that doesn't require partial rewards

[mtrofin]

Summary

If we use training algorithms that don't need partial rewards, we don't
need to worry about an ir2native model. In that case, training logs
won't contain a 'delta_size' feature either (since that's the partial
reward).

[yundiqian]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

458

I'm not familiar with the UI. Is the '>>' expected?

[MaskRay]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

458

It is expected. If Phabricator detects indentation changes, it will print >>

llvm/test/Transforms/Inline/ML/development-training-log.ll

4

--check-prefixes

[mtrofin]

feedback

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

DevelopmentModeMLInlineAdvisor(
       Module &M, ModuleAnalysisManager &MAM,
       std::unique_ptr<MLModelRunner> ModelRunner,
       std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference,
       std::unique_ptr<TrainingLogger> Logger);
 
   size_t getTotalSizeEstimate();
 
   virtual ~DevelopmentModeMLInlineAdvisor();
-  void updateNativeSizeEstimate(int64_t Change) { CurrentNativeSize += Change; }
+  void updateNativeSizeEstimate(int64_t Change) {
+    *CurrentNativeSize += Change;
+  }
   void resetNativeSize(Function *F) {
     FAM.invalidate<InlineSizeEstimatorAnalysis>(*F);
   }
 
   std::unique_ptr<MLInlineAdvice>
   getMandatoryAdvice(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
   std::unique_ptr<MLInlineAdvice>
   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
 
-  size_t getNativeSizeEstimate(const Function &F) const;
+  Optional<size_t> getNativeSizeEstimate(const Function &F) const;
 
 private:
   bool isLogging() const { return !!Logger; }
 
   std::function<bool(CallBase &)> GetDefaultAdvice;
   const bool IsDoingInference;
   std::unique_ptr<TrainingLogger> Logger;
 
-  const int32_t InitialNativeSize;
-  int32_t CurrentNativeSize = 0;
+  const Optional<int32_t> InitialNativeSize;
+  Optional<int32_t> CurrentNativeSize;
 };
 
 /// A variant of MLInlineAdvice that tracks all non-trivial inlining
 /// decisions, for training/logging.
 class LoggingMLInlineAdvice : public MLInlineAdvice {
 public:
   LoggingMLInlineAdvice(DevelopmentModeMLInlineAdvisor *Advisor, CallBase &CB,
                         OptimizationRemarkEmitter &ORE, bool Recommendation,
-                        TrainingLogger &Logger, size_t CallerSizeEstimateBefore,
-                        size_t CalleeSizeEstimateBefore, bool DefaultDecision,
-                        bool Mandatory = false)
+                        TrainingLogger &Logger,
+                        Optional<size_t> CallerSizeEstimateBefore,
+                        Optional<size_t> CalleeSizeEstimateBefore,
+                        bool DefaultDecision, bool Mandatory = false)
       : MLInlineAdvice(Advisor, CB, ORE, Recommendation), Logger(Logger),
         CallerSizeEstimateBefore(CallerSizeEstimateBefore),
         CalleeSizeEstimateBefore(CalleeSizeEstimateBefore),
         DefaultDecision(DefaultDecision), Mandatory(Mandatory) {}
 
   virtual ~LoggingMLInlineAdvice() = default;
 
 private:
   DevelopmentModeMLInlineAdvisor *getAdvisor() const {
     return static_cast<DevelopmentModeMLInlineAdvisor *>(Advisor);
   }
   void recordInliningImpl() override {
     MLInlineAdvice::recordInliningImpl();
     getAdvisor()->resetNativeSize(Caller);
     int Reward = std::numeric_limits<int>::max();
-    if (!getAdvisor()->isForcedToStop()) {
-      int NativeSizeAfter = getAdvisor()->getNativeSizeEstimate(*Caller) +
-                            CalleeSizeEstimateBefore;
+    if (InlineSizeEstimatorAnalysis::isEvaluatorRequested() &&
+        !getAdvisor()->isForcedToStop()) {
+      int NativeSizeAfter = *getAdvisor()->getNativeSizeEstimate(*Caller) +
+                            *CalleeSizeEstimateBefore;
       Reward = NativeSizeAfter -
-               (CallerSizeEstimateBefore + CalleeSizeEstimateBefore);
+               (*CallerSizeEstimateBefore + *CalleeSizeEstimateBefore);
       getAdvisor()->updateNativeSizeEstimate(Reward);
     }
     log(Reward, /*Success=*/true);
   }
 
   void recordInliningWithCalleeDeletedImpl() override {
     MLInlineAdvice::recordInliningWithCalleeDeletedImpl();
     getAdvisor()->resetNativeSize(Caller);
-    if (!getAdvisor()->isForcedToStop()) {
-      int NativeSizeAfter = getAdvisor()->getNativeSizeEstimate(*Caller);
+    if (InlineSizeEstimatorAnalysis::isEvaluatorRequested() &&
+        !getAdvisor()->isForcedToStop()) {
+      int NativeSizeAfter = *getAdvisor()->getNativeSizeEstimate(*Caller);
       int Reward = NativeSizeAfter -
-                   (CallerSizeEstimateBefore + CalleeSizeEstimateBefore);
+                   (*CallerSizeEstimateBefore + *CalleeSizeEstimateBefore);
       getAdvisor()->updateNativeSizeEstimate(Reward);
       log(Reward, /*Success=*/true);
     }
   }
 
   void recordUnsuccessfulInliningImpl(const InlineResult &Result) override {
     MLInlineAdvice::recordUnsuccessfulInliningImpl(Result);
     log(NoReward, /*Success=*/false);
void log(int64_t Reward, bool Success) {
     Event.DefaultDecision = DefaultDecision;
     Event.Effect = Success;
     Event.Reward = Reward;
     Logger.logInlineEvent(Event, getAdvisor()->getModelRunner());
   }
 
   static const int64_t NoReward = 0;
   TrainingLogger &Logger;
-  const size_t CallerSizeEstimateBefore;
-  const size_t CalleeSizeEstimateBefore;
+  const Optional<size_t> CallerSizeEstimateBefore;
+  const Optional<size_t> CalleeSizeEstimateBefore;
   const bool DefaultDecision;
   const bool Mandatory;
 };
 
 /// A pseudo model runner. We use it to store feature values when collecting
 /// logs for the default policy, but never ask it to 'run'.
 class NoInferenceModelRunner : public MLModelRunner {
 public:
void TrainingLogger::print() {
 
   writeTensorsAsFeatureLists(
       OutFile, TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}),
       DefaultDecisions.data(), NumberOfRecords);
 
   writeRawTensorsAsFeatureLists(
       OutFile, TensorSpec::createSpec<int64_t>(DecisionName, {1}),
       Outputs[0].data(), NumberOfRecords);
+
+  if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
     writeTensorsAsFeatureLists(OutFile,
                                TensorSpec::createSpec<int64_t>(RewardName, {1}),
                                Rewards.data(), NumberOfRecords);
 
   for (size_t I = 1; I < Outputs.size(); ++I)
     writeRawTensorsAsFeatureLists(OutFile, MUTR->outputSpecs()[I],
                                   Outputs[I].data(), NumberOfRecords,
                                   StringRef(MUTR->outputNames()[I]));
 
   OutFile << "}\n";
 }
DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
   assert(IsDoingInference || isLogging());
 }
 
 DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
   if (isLogging())
     Logger->print();
 }
 
-size_t
+Optional<size_t>
 DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
+  if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
+    return None;
   auto &R =
       FAM.getResult<InlineSizeEstimatorAnalysis>(const_cast<Function &>(F));
   if (!R) {
     F.getParent()->getContext().emitError(
         "Native size estimator is not present.");
     return 0;
   }
   return *R;
 }
 
 std::unique_ptr<MLInlineAdvice>
 DevelopmentModeMLInlineAdvisor::getMandatoryAdvice(
     CallBase &CB, OptimizationRemarkEmitter &ORE) {
   if (!isLogging())
     return MLInlineAdvisor::getMandatoryAdvice(CB, ORE);
+
   return std::make_unique<LoggingMLInlineAdvice>(
       /*Advisor=*/this,
       /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/true, /*Logger=*/*Logger,
       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
       /*CalleeSizeEstimateBefore=*/
       getNativeSizeEstimate(*CB.getCalledFunction()),
       /*DefaultDecision=*/true, /*Mandatory*/ true);
 }
return std::make_unique<LoggingMLInlineAdvice>(
       /*Logger=*/*Logger,
       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
       /*CalleeSizeEstimateBefore=*/
       getNativeSizeEstimate(*CB.getCalledFunction()),
       /*DefaultDecision=*/DefaultAdvice);
 }
 
 size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() {
+  if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
+    return 0;
   size_t Ret = 0;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     if (isFunctionDeleted(&F))
       continue;
-    Ret += getNativeSizeEstimate(F);
+    Ret += *getNativeSizeEstimate(F);
   }
   return Ret;
 }
 
 ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
                                                    const std::string &ModelPath)
     : MLModelRunner(Ctx) {
   std::vector<TensorSpec> InputSpecs;
void ModelUnderTrainingRunner::setFeature(FeatureIndex Index, int64_t Value) {
   size_t NumericIndex = static_cast<size_t>(Index);
   *(Evaluator->getInput<int64_t>(NumericIndex)) = Value;
 }
 
 std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
     Module &M, ModuleAnalysisManager &MAM,
     std::function<bool(CallBase &)> GetDefaultAdvice) {
   auto &Ctx = M.getContext();
-  if (TrainingLog.empty() !=
-      !InlineSizeEstimatorAnalysis::isEvaluatorRequested()) {
-    Ctx.emitError("For development mode, if training logs are requested, then "
-                  "a size estimator must be available; either that, or neither "
-                  "are specified.");
-    return nullptr;
-  }
-
   std::unique_ptr<MLModelRunner> Runner;
   ModelUnderTrainingRunner *MUTRPtr = nullptr;
   bool IsDoingInference = false;
   if (TFModelUnderTrainingPath.empty())
     Runner.reset(new NoInferenceModelRunner(Ctx));
   else {
     auto MUTR = std::make_unique<ModelUnderTrainingRunner>(
         Ctx, TFModelUnderTrainingPath);

llvm/test/Transforms/Inline/ML/development-training-log.ll

 ; Test that we can produce a log if we have or do not have a model, in development mode.
 ; REQUIRES: have_tf_api
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s 
-; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-output-spec-override=%S/Inputs/test_output_spec.json -S < %s | FileCheck %s --check-prefix=EXTRA-OUTPUTS
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-output-spec-override=%S/Inputs/test_output_spec.json -S < %s | FileCheck %s --check-prefix=EXTRA-OUTPUTS --check-prefix=CHECK
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -S < %s | FileCheck %s --check-prefix=NOREWARD
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -S < %s | FileCheck %s --check-prefix=NOREWARD
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-linux-gnu"
 
 declare i32 @f1(i32)
 declare i32 @f2(i32)
 
 define dso_local i32 @branches(i32) {

 ; CHECK-NEXT:       feature: { int64_list: { value: [1] } }
 ; CHECK-NEXT:     }
 ; CHECK-NEXT:   }
 ; CHECK-NEXT:   feature_list: {
 ; CHECK-NEXT:     key: "delta_size" value: {
 ; CHECK-NEXT:       feature: { int64_list: { value: [0] } }
 ; CHECK-NEXT:     }
 ; CHECK-NEXT:   }
+; NOREWARD-NOT: key: "delta_size" value: {
 ; CHECK-NOT: fake_extra_output
 ; EXTRA-OUTPUTS:          key: "fake_extra_output" value: {
 ; EXTRA-OUTPUTS-NEXT:       feature: { int64_list: { value: [1] } }
 ; EXTRA-OUTPUTS-NEXT:     }
 ; EXTRA-OUTPUTS-NEXT:   }
 No newline at end of file