[llvm][MLInliner] Don't log 'mandatory' events

[mtrofin]

Summary

We don't want mandatory events in the training log. We do want to handle
them, to keep the native size accounting accurate, but that's all.

Fixed the code, also expanded the test to capture this.

[MaskRay]

llvm/test/Transforms/Inline/ML/bounds-checks-rewards.ll

21

Note that %S/../../ is uncommon. This will break our internal lit runner.

[mtrofin]

llvm/test/Transforms/Inline/ML/bounds-checks-rewards.ll

21

Ack - as discussed offline, I would also prefer moving the inputs to the local Inputs dir. For the interim (i.e. until we actually enable this by default in our internal build, for example, or until we move the files), I'll add a comment warning test maintainers about the cross-folder dependency.

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

private:
   void writeTensor(raw_fd_ostream &OutFile, StringRef TensorName,
                    const std::vector<T> &Tensor);
 
   std::vector<InlineFeatures> Features;
   std::vector<bool> DefaultDecisions;
   std::vector<bool> Decisions;
   std::vector<bool> Effects;
   std::vector<int64_t> Rewards;
-  std::vector<bool> Mandatory;
 };
 
 /// An extension of the MLInlineAdvisor for the 'development' mode, targeting
 /// the offline training scenario. Note that training happens outside of the
 /// compiler, this facility is concerned with producing training data ("logs").
 /// This InlineAdvisor can operate in the following modes:
 ///
 /// 1) collect logs for the default policy. This is useful for bootstrapping

 
 /// 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)
+                        size_t CalleeSizeEstimateBefore, bool DefaultDecision,
+                        bool Mandatory = false)
       : MLInlineAdvice(Advisor, CB, ORE, Recommendation), Logger(Logger),
         CallerSizeEstimateBefore(CallerSizeEstimateBefore),
         CalleeSizeEstimateBefore(CalleeSizeEstimateBefore),
-        DefaultDecision(DefaultDecision) {}
+        DefaultDecision(DefaultDecision), Mandatory(Mandatory) {}
 
   virtual ~LoggingMLInlineAdvice() = default;
 
 private:
   DevelopmentModeMLInlineAdvisor *getAdvisor() const {
     return static_cast<DevelopmentModeMLInlineAdvisor *>(Advisor);
   }
   void recordInliningImpl() override {
private:
   }
 
   void recordUnattemptedInliningImpl() override {
     MLInlineAdvice::recordUnattemptedInliningImpl();
     log(NoReward, /*Success=*/false);
   }
 
   void log(int64_t Reward, bool Success) {
+    if (Mandatory)
+      return;
     InlineEvent Event;
     Event.AdvisedDecision = isInliningRecommended();
     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 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:
   NoInferenceModelRunner(LLVMContext &Ctx)
       : MLModelRunner(Ctx), Features(NumberOfFeatures) {}
DevelopmentModeMLInlineAdvisor::getMandatoryAdvice(
   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);
+      /*DefaultDecision=*/true, /*Mandatory*/ true);
 }
 
 std::unique_ptr<MLInlineAdvice>
 DevelopmentModeMLInlineAdvisor::getAdviceFromModel(
     CallBase &CB, OptimizationRemarkEmitter &ORE) {
   if (IsDoingInference && !isLogging())
     return MLInlineAdvisor::getAdviceFromModel(CB, ORE);
 

llvm/test/Transforms/Inline/ML/bounds-checks-rewards.ll

 ; Test behavior when inlining policy grows size out of control.
 ; In all cases, the end result is the same: mandatory inlinings must happen.
+; Also in all cases, we don't record the mandatory inlining (there's nothing to
+; learn from it).
 ; However, when we discover we 'trip' over the artificially-low size increase 
 ; factor, we penalize the 'bad' decision.
 ; REQUIRES: have_tf_api
+;
+; When the bounds are very wide ("no bounds"), all inlinings happen.
 ; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -training-log=- -enable-ml-inliner=development -ml-advisor-size-increase-threshold=10.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=NOBOUNDS
+;
+; When the bounds are very restrictive, the first inlining happens but it's
+; considered as "bad" (since it trips over the bounds) and its reward is a
+; penalty. However, the mandatory inlining, which is considered next, happens.
+; No other inlinings happend.
 ; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -training-log=- -enable-ml-inliner=development -ml-advisor-size-increase-threshold=1.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=BOUNDS
+;
+; With more restrictive bounds, the first inlining happens and is OK. The
+; mandatory inlining happens next, and it trips over the bounds, which then
+; forces no further inlinings.
+; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%S/../../../../lib/Analysis/models/inliner -training-log=- -enable-ml-inliner=development -ml-advisor-size-increase-threshold=1.1 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=RELAXED-BOUNDS
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-grtev4-linux-gnu"
 
 declare i64 @f1()
 
-define i64 @f2() #0 {
+define i64 @may_not_be_inlined() {
   %r = call i64 @f1()
   %r2 = add i64 13, %r
   ret i64 %r2
 }
 
-define i64 @some_function() {
-  %r = call i64 @f1()
+define i64 @must_be_inlined() #0 {
+  %r = call i64 @may_not_be_inlined()
   %r2 = add i64 13, %r
   ret i64 %r2
 }
 
 define i64 @top() {
-  %r = call i64 @f2()
-  %r2 = call i64 @some_function()
-  %r3 = add i64 %r, %r2
-  ret i64 %r3
+  %r = call i64 @must_be_inlined()
+  %r2 = call i64 @may_not_be_inlined()
+  %r3 = call i64 @may_not_be_inlined()
+  %r4 = add i64 %r, %r2
+  %r5 = add i64 %r3, %r4
+  ret i64 %r5
 }
 
 attributes #0 = { alwaysinline }
 ; CHECK: key: "delta_size" value: {
-; NOBOUNDS-NEXT: feature: { int64_list: { value: [10] } }
 ; NOBOUNDS-NEXT: feature: { int64_list: { value: [6] } }
+; RELAXED-BOUNDS-NEXT: feature: { int64_list: { value: [6] } }
+; NOBOUNDS-NEXT: feature: { int64_list: { value: [-11] } }
+; NOBOUNDS-NEXT: feature: { int64_list: { value: [4] } }
 ; BOUNDS-NEXT: feature: { int64_list: { value: [2147483647] } }
 ; CHECK-NEXT: }
 ; CHECK-LABEL: @top
-; f2 must always be inlined, so we won't find a call to it in @top()
-; CHECK-NOT: call i64 @f2
+; must_be_inlined must always be inlined, so we won't find a call to it in @top()
+; CHECK-NOT: call i64 @must_be_inlined
 ; @some-function isn't mandatory, and when we set the increase threshold too low,
 ; it won't be inlined.
-; NOBOUNDS-NOT: @some_function
-; BOUNDS: call i64 @some_function
+; NOBOUNDS-NOT: @may_not_be_inlined
+; RELAXED-BOUNDS: call i64 @may_not_be_inlined
+; BOUNDS: call i64 @may_not_be_inlined
 No newline at end of file