[mlgo] Hook up the interactive runner to the mlgo-ed passes

[mtrofin]

Summary

This hooks up the interactive model runner to the passes that support
ml-based decisions. Because the interface to this runner is the exact
same as the one used during inference, we just reuse the exact same
setup we have for "release mode". This makes "release mode" a misnomer -
and that's something we needed to resolve sooner or later (e.g.
supporting more than one embedded model for the same problem was another
reason to drop that nomenclature). That will happen in a subsequent
change.

To use this evaluator, just enable the pass in (currently) "release"
mode, but also pass the base name for the 2 channel files via the
pass-specific flag.

The 2 files are the responsibilty of the hosting process. The added
tests use a minimal, toy such host, illustrating setup and
communication.

[kazu]

llvm/lib/Analysis/models/interactive_host.py

2

Could you add some comments about what this script is for and how to use it (or a pointer to a .ll test that uses it)?

IIUC, the script invokes the compiler and acts as a bridge between the compiler proper and the user's model. The user is responsible for supplying the complete command line, including the compiler proper (clang, opt, etc) and the -regalloc-priority-interactive-channel-base= option to specify the base name (that is, without .in or .out) of the named pipes.

llvm/test/Transforms/Inline/ML/Inputs/interactive_main.py

12

Could you add a comment somewhere in this file?

IIUC, we say "yes, no, yes, no, ..." as expected by llvm/test/Transforms/Inline/ML/interactive-mode.ll.

[mtrofin]

feedback

llvm/lib/Analysis/models/interactive_host.py

2

done - expanded a bit on the args, too.

[kazu]

Thanks for the update!

[chapuni]

llvm/test/CodeGen/MLRegalloc/Inputs/interactive_main.py

12

Seems unsupported in py-3.8. Tweaked in e006c7dfa79a

[dyung]

There still seems to be failures related to this change, can you take a look please?

https://lab.llvm.org/buildbot/#/builders/247/builds/1260

Traceback (most recent call last):
  File "/home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/build/test/Transforms/Inline/ML/Output/interactive-mode.ll.tmp.rundir/interactive_main.py", line 21, in <module>
    main(sys.argv[1:])
  File "/home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/build/test/Transforms/Inline/ML/Output/interactive-mode.ll.tmp.rundir/interactive_main.py", line 17, in main
    interactive_host.run_interactive(args[0], a.advice, args[1:])
  File "/home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/build/test/Transforms/Inline/ML/Output/interactive-mode.ll.tmp.rundir/interactive_host.py", line 71, in run_interactive
    context, next_event, fc, tensor_specs, None)
  File "/home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/build/test/Transforms/Inline/ML/Output/interactive-mode.ll.tmp.rundir/log_reader.py", line 100, in read_one_observation
    features.append(read_tensor(f, ts))
  File "/home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/build/test/Transforms/Inline/ML/Output/interactive-mode.ll.tmp.rundir/log_reader.py", line 71, in read_tensor
    size = math.prod(ts.shape) * ctypes.sizeof(ts.element_type)
AttributeError: module 'math' has no attribute 'prod'
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/build/bin/FileCheck /home/buildbot/buildbot-root/llvm-clang-x86_64-gcc-ubuntu/llvm-project/llvm/test/Transforms/Inline/ML/interactive-mode.ll

[mtrofin]

In D143218#4104421, @dyung wrote:

There still seems to be failures related to this change, can you take a look please?

https://lab.llvm.org/buildbot/#/builders/247/builds/1260

Fixed in 062380c86f061c7308c8b973a1b94e1f6323be5d.

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

 constexpr FeatureIndex
 inlineCostFeatureToMlFeature(InlineCostFeatureIndex Feature) {
   return static_cast<FeatureIndex>(static_cast<size_t>(Feature));
 }
 
 constexpr size_t NumberOfFeatures =
     static_cast<size_t>(FeatureIndex::NumberOfFeatures);
 
-extern const std::array<TensorSpec, NumberOfFeatures> FeatureMap;
+extern const std::vector<TensorSpec> FeatureMap;
 
 extern const char *const DecisionName;
+extern const TensorSpec InlineDecisionSpec;
 extern const char *const DefaultDecisionName;
 extern const char *const RewardName;
 
 using InlineFeatures = std::vector<int64_t>;
 
 } // namespace llvm
 #endif // LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H

llvm/include/llvm/Analysis/InteractiveModelRunner.h

public:
   InteractiveModelRunner(LLVMContext &Ctx,
                          const std::vector<TensorSpec> &Inputs,
                          const TensorSpec &Advice, StringRef OutboundName,
                          StringRef InboundName);
 
   static bool classof(const MLModelRunner *R) {
     return R->getKind() == MLModelRunner::Kind::Interactive;
   }
-  void switchContext(StringRef Name) {
+  void switchContext(StringRef Name) override {
     Log->switchContext(Name);
     Log->flush();
   }
 
   virtual ~InteractiveModelRunner();
 
 private:
   void *evaluateUntyped() override;

llvm/include/llvm/Analysis/MLModelRunner.h

public:
 
   void *getTensorUntyped(size_t Index) { return InputBuffers[Index]; }
   const void *getTensorUntyped(size_t Index) const {
     return (const_cast<MLModelRunner *>(this))->getTensorUntyped(Index);
   }
 
   enum class Kind : int { Unknown, Release, Development, NoOp, Interactive };
   Kind getKind() const { return Type; }
+  virtual void switchContext(StringRef Name) {}
 
 protected:
   MLModelRunner(LLVMContext &Ctx, Kind Type, size_t NrInputs)
       : Ctx(Ctx), Type(Type), InputBuffers(NrInputs) {
     assert(Type != Kind::Unknown);
   }
   virtual void *evaluateUntyped() = 0;
 

llvm/include/llvm/Analysis/ReleaseModeModelRunner.h

public:
   NoopSavedModelImpl() = default;
   int LookupArgIndex(const std::string &) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
   int LookupResultIndex(const std::string &) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
   void Run() { llvm_unreachable(NOOP_MODEL_ERRMSG); }
   void *result_data(int) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
   void *arg_data(int) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
 #undef NOOP_MODEL_ERRMSG
 };
+
+template <class T> bool isEmbeddedModelEvaluatorValid() { return true; }
+
+template <> inline bool isEmbeddedModelEvaluatorValid<NoopSavedModelImpl>() {
+  return false;
+}
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

TrainingLogger::TrainingLogger(StringRef LogFileName,
 
   if (MUTR)
     append_range(FT, MUTR->extraOutputsForLoggingSpecs());
 
   DefaultDecisionPos = FT.size();
   FT.push_back(TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}));
 
   DecisionPos = FT.size();
-  FT.push_back(TensorSpec::createSpec<int64_t>(DecisionName, {1}));
+  FT.push_back(InlineDecisionSpec);
   std::error_code EC;
   auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
   if (EC)
     dbgs() << (EC.message() + ":" + TrainingLog);
 
   L = std::make_unique<Logger>(
       std::move(OS), FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
       InlineSizeEstimatorAnalysis::isEvaluatorRequested());

llvm/lib/Analysis/InlineAdvisor.cpp

#ifdef LLVM_HAVE_TFLITE
     Advisor =
         llvm::getDevelopmentModeAdvisor(M, MAM, [&FAM, Params](CallBase &CB) {
           auto OIC = getDefaultInlineAdvice(CB, FAM, Params);
           return OIC.has_value();
         });
 #endif
     break;
   case InliningAdvisorMode::Release:
-#ifdef LLVM_HAVE_TF_AOT
     LLVM_DEBUG(dbgs() << "Using release-mode inliner policy.\n");
     Advisor = llvm::getReleaseModeAdvisor(M, MAM);
-#endif
     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

 //===----------------------------------------------------------------------===//
 #include "llvm/Analysis/MLInlineAdvisor.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/InlineModelFeatureMaps.h"
+#include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/Analysis/ReleaseModeModelRunner.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 
 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"));
+
 #if defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
-#include "llvm/Analysis/ReleaseModeModelRunner.h"
 // 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) {
-  auto AOTRunner =
-      std::make_unique<ReleaseModeModelRunner<llvm::InlinerSizeModel>>(
-          M.getContext(), FeatureMap, DecisionName);
+  if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
+      InteractiveChannelBaseName.empty())
+    return nullptr;
+  std::unique_ptr<MLModelRunner> AOTRunner;
+  if (InteractiveChannelBaseName.empty())
+    AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
+  else
+    AOTRunner = std::make_unique<InteractiveModelRunner>(
+        M.getContext(), FeatureMap, InlineDecisionSpec,
+        InteractiveChannelBaseName + ".out",
+        InteractiveChannelBaseName + ".in");
   return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
 }
-#endif
 
 #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."),
     cl::init(2.0));
 
 static cl::opt<bool> KeepFPICache(
     "ml-advisor-keep-fpi-cache", cl::Hidden,
     cl::desc(
         "For test - keep the ML Inline advisor's FunctionPropertiesInfo cache"),
     cl::init(false));
 
 // clang-format off
-const std::array<TensorSpec, NumberOfFeatures> llvm::FeatureMap{
+const std::vector<TensorSpec> llvm::FeatureMap{
 #define POPULATE_NAMES(_, NAME) TensorSpec::createSpec<int64_t>(NAME, {1} ),
 // InlineCost features - these must come first
   INLINE_COST_FEATURE_ITERATOR(POPULATE_NAMES)
 #undef POPULATE_NAMES
 
 // Non-cost features
 #define POPULATE_NAMES(_, NAME, __) TensorSpec::createSpec<int64_t>(NAME, {1} ),
   INLINE_FEATURE_ITERATOR(POPULATE_NAMES)
 #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 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)
     : InlineAdvisor(
           M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
       ModelRunner(std::move(Runner)),
       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
   // heuristic's decisions - and, thus, equally important for training for
   // improvement.
   CallGraph CGraph(M);
   for (auto I = scc_begin(&CGraph); !I.isAtEnd(); ++I) {

llvm/lib/Analysis/models/interactive_host.py

+"""Utility for testing InteractiveModelRunner.
+
+Use it from pass-specific tests by providing a main .py which calls this library's
+`run_interactive` with an appropriate callback to provide advice.
+
+From .ll tests, just call the above-mentioned main as a prefix to the opt/llc
+invocation (with the appropriate flags enabling the interactive mode)
+
+Examples:
+test/Transforms/Inline/ML/interactive-mode.ll
+test/CodeGen/MLRegalloc/interactive-mode.ll
+"""
+
+import ctypes
+import log_reader
+import io
+import math
+import os
+import subprocess
+from typing import BinaryIO, Callable
+
+
+def send(f: io.BufferedWriter, value: int | float, spec: log_reader.TensorSpec):
+  """Send the `value` - currently just a scalar - formatted as per `spec`."""
+
+  # just int64 for now
+  assert (spec.element_type == ctypes.c_int64)
+  to_send = ctypes.c_int64(int(value))
+  assert f.write(bytes(to_send)) == ctypes.sizeof(
+      spec.element_type) * math.prod(spec.shape)
+  f.flush()
+
+
+def run_interactive(temp_rootname: str,
+                    make_response: Callable[[list[log_reader.TensorValue]],
+                                            int | float],
+                    process_and_args: list[str]):
+  """Host the compiler.
+  Args:
+    temp_rootname: the base file name from which to construct the 2 pipes for
+    communicating with the compiler.
+    make_response: a function that, given the current tensor values, provides a
+    response.
+    process_and_args: the full commandline for the compiler. It it assumed it
+    contains a flag poiting to `temp_rootname` so that the InteractiveModeRunner
+    would attempt communication on the same pair as this function opens.
+
+  This function sets up the communication with the compiler - via 2 files named
+  `temp_rootname`.in and `temp_rootname`.out - prints out the received features,
+  and sends back to the compiler an advice (which it gets from `make_response`).
+  It's used for testing, and also to showcase how to set up communication in an
+  interactive ML ("gym") environment.
+  """
+  to_compiler = temp_rootname + ".in"
+  from_compiler = temp_rootname + ".out"
+  try:
+    os.mkfifo(to_compiler, 0o666)
+    os.mkfifo(from_compiler, 0o666)
+    compiler_proc = subprocess.Popen(process_and_args)
+    with io.BufferedWriter(io.FileIO(to_compiler, 'wb')) as tc:
+      with io.BufferedReader(io.FileIO(from_compiler, 'rb')) as fc:
+        tensor_specs, _, advice_spec = log_reader.read_header(fc)
+        context = None
+        while compiler_proc.poll() is None and (next_event := fc.readline()):
+          last_context, observation_id, features, _ = log_reader.read_one_observation(
+              context, next_event, fc, tensor_specs, None)
+          if last_context != context:
+            print(f'context: {last_context}')
+          context = last_context
+          print(f'observation: {observation_id}')
+          tensor_values = []
+          for fv in features:
+            log_reader.pretty_print_tensor_value(fv)
+            tensor_values.append(fv)
+          send(tc, make_response(tensor_values), advice_spec)
+
+  finally:
+    os.unlink(to_compiler)
+    os.unlink(from_compiler)

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 //===- MLRegAllocEvictAdvisor.cpp - ML eviction advisor -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of the ML eviction advisor and reward injection pass
 //
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
 #include "RegAllocEvictionAdvisor.h"
 #include "RegAllocGreedy.h"
+#include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TFLITE)
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
 #include "llvm/Analysis/Utils/TrainingLogger.h"
 #endif
 #include "MLRegallocEvictAdvisor.h"

 // Generated header in release (AOT) mode
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
 #include "RegallocEvictModel.h"
 using CompiledModelType = RegallocEvictModel;
 #else
 using CompiledModelType = NoopSavedModelImpl;
 #endif
 
+static cl::opt<std::string> InteractiveChannelBaseName(
+    "regalloc-evict-interactive-channel-base", cl::Hidden,
+    cl::desc(
+        "Base file path for the interactive mode. The incoming filename should "
+        "have the name <regalloc-evict-interactive-channel-base>.in, while the "
+        "outgoing name should be "
+        "<regalloc-evict-interactive-channel-base>.out"));
+
 // Options that only make sense in development mode
 #ifdef LLVM_HAVE_TFLITE
 #include "RegAllocScore.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 
 static cl::opt<std::string> TrainingLog(
     "regalloc-training-log", cl::Hidden,
     cl::desc("Training log for the register allocator eviction model"));

 #define RA_EVICT_REST_DEVELOPMENT_FEATURES(M)
 #endif
 
 // The model learns to pick one of the mask == 1 interferences. This is the
 // name of the output tensor. The contract with the model is that the output
 // will be guaranteed to be to a mask == 1 position. Using a macro here to
 // avoid 'not used' warnings (and keep cond compilation to a minimum)
 #define DecisionName "index_to_evict"
+static const TensorSpec DecisionSpec =
+    TensorSpec::createSpec<int64_t>(DecisionName, {1});
 
 // Named features index.
 enum FeatureIDs {
 #define _FEATURE_IDX_SIMPLE(_, name, __, ___) name
 #define _FEATURE_IDX(A, B, C, D) _FEATURE_IDX_SIMPLE(A, B, C, D),
   RA_EVICT_FEATURES_LIST(_FEATURE_IDX) FeatureCount,
 #ifdef LLVM_HAVE_TFLITE
   RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_FEATURE_IDX_SIMPLE) = FeatureCount,
private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineLoopInfo>();
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    if (!Runner)
+    if (!Runner) {
+      if (InteractiveChannelBaseName.empty())
         Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
             MF.getFunction().getContext(), InputFeatures, DecisionName);
+      else
+        Runner = std::make_unique<InteractiveModelRunner>(
+            MF.getFunction().getContext(), InputFeatures, DecisionSpec,
+            InteractiveChannelBaseName + ".out",
+            InteractiveChannelBaseName + ".in");
+    }
     return std::make_unique<MLEvictAdvisor>(
         MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
         getAnalysis<MachineLoopInfo>());
   }
-  std::unique_ptr<ReleaseModeModelRunner<CompiledModelType>> Runner;
+  std::unique_ptr<MLModelRunner> Runner;
 };
 
 // ===================================
 // Development mode-specifics
 // ===================================
 //
 // Features we log
 #ifdef LLVM_HAVE_TFLITE
-static const TensorSpec Output =
-    TensorSpec::createSpec<int64_t>(DecisionName, {1});
 static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
 
 // Features we bind on the model. The tensor names have a prefix, and we also
 // need to include some tensors that are expected to be present by the
 // training algo.
 // TODO: can we just get rid of these?
 #define _DECL_TRAIN_FEATURES(type, name, shape, _)                             \
   TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
if (EC) {
       return false;
     }
     std::vector<TensorSpec> LFS = InputFeatures;
     if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
       append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
     // We always log the output; in particular, if we're not evaluating, we
     // don't have an output spec json file. That's why we handle the
     // 'normal' output separately.
-    LFS.push_back(Output);
+    LFS.push_back(DecisionSpec);
 
     Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
                                    /*IncludeReward*/ true);
     return false;
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {

 MLEvictAdvisor::MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
                                MLModelRunner *Runner,
                                const MachineBlockFrequencyInfo &MBFI,
                                const MachineLoopInfo &Loops)
     : RegAllocEvictionAdvisor(MF, RA), DefaultAdvisor(MF, RA),
       Runner(std::move(Runner)), MBFI(MBFI), Loops(Loops),
       InitialQSize(MLEvictAdvisor::getInitialQueueSize(MF)) {
   assert(this->Runner);
+  Runner->switchContext(MF.getName());
   DoNotNormalize.set(FeatureIDs::mask);
   DoNotNormalize.set(FeatureIDs::is_free);
   DoNotNormalize.set(FeatureIDs::is_hint);
   DoNotNormalize.set(FeatureIDs::is_local);
   DoNotNormalize.set(FeatureIDs::min_stage);
   DoNotNormalize.set(FeatureIDs::max_stage);
   DoNotNormalize.set(FeatureIDs::progress);
 }
getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF,
                                                                    GetReward);
   getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF,
                                                                    GetReward);
   return false;
 }
 #endif // #ifdef LLVM_HAVE_TFLITE
 
 RegAllocEvictionAdvisorAnalysis *llvm::createReleaseModeAdvisor() {
-  return new ReleaseModeEvictionAdvisorAnalysis();
+  return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
+                 !InteractiveChannelBaseName.empty()
+             ? new ReleaseModeEvictionAdvisorAnalysis()
+             : nullptr;
 }
 
 // In all cases except development mode, we don't need scoring.
 #if !defined(LLVM_HAVE_TFLITE)
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &) { return false; }
 #endif

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

 //===- MLRegAllocPriorityAdvisor.cpp - ML priority advisor-----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of the ML priority advisor and reward injection pass
 //
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
 #include "RegAllocGreedy.h"
 #include "RegAllocPriorityAdvisor.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/InteractiveModelRunner.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/ReleaseModeModelRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/CodeGen/CalcSpillWeights.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"

 #if defined(LLVM_HAVE_TFLITE)
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
 #include "llvm/Analysis/Utils/TrainingLogger.h"
 #endif
 
 using namespace llvm;
 
+static cl::opt<std::string> InteractiveChannelBaseName(
+    "regalloc-priority-interactive-channel-base", cl::Hidden,
+    cl::desc(
+        "Base file path for the interactive mode. The incoming filename should "
+        "have the name <regalloc-priority-interactive-channel-base>.in, while "
+        "the outgoing name should be "
+        "<regalloc-priority-interactive-channel-base>.out"));
+
+using CompiledModelType = NoopSavedModelImpl;
+
 // Options that only make sense in development mode
 #ifdef LLVM_HAVE_TFLITE
 #include "RegAllocScore.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 
 static cl::opt<std::string> TrainingLog(
     "regalloc-priority-training-log", cl::Hidden,
     cl::desc("Training log for the register allocator priority model"));

 static const std::vector<int64_t> PerLiveRangeShape{1};
 
 #define RA_PRIORITY_FEATURES_LIST(M)                                           \
   M(int64_t, li_size, PerLiveRangeShape, "size")                               \
   M(int64_t, stage, PerLiveRangeShape, "stage")                                \
   M(float, weight, PerLiveRangeShape, "weight")
 
 #define DecisionName "priority"
+static const TensorSpec DecisionSpec =
+    TensorSpec::createSpec<float>(DecisionName, {1});
+
 
 // Named features index.
 enum FeatureIDs {
 #define _FEATURE_IDX(_, name, __, ___) name,
   RA_PRIORITY_FEATURES_LIST(_FEATURE_IDX)
 #undef _FEATURE_IDX
       FeatureCount
 };
private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     AU.addRequired<SlotIndexes>();
     RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    if (!Runner)
-      Runner = std::make_unique<ReleaseModeModelRunner<NoopSavedModelImpl>>(
+    if (!Runner) {
+      if (InteractiveChannelBaseName.empty())
+        Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
             MF.getFunction().getContext(), InputFeatures, DecisionName);
+      else
+        Runner = std::make_unique<InteractiveModelRunner>(
+            MF.getFunction().getContext(), InputFeatures, DecisionSpec,
+            InteractiveChannelBaseName + ".out",
+            InteractiveChannelBaseName + ".in");
+    }
     return std::make_unique<MLPriorityAdvisor>(
         MF, RA, &getAnalysis<SlotIndexes>(), Runner.get());
   }
-  std::unique_ptr<ReleaseModeModelRunner<NoopSavedModelImpl>> Runner;
+  std::unique_ptr<MLModelRunner> Runner;
 };
 
 // ===================================
 // Development mode-specifics
 // ===================================
 //
 // Features we log
 #ifdef LLVM_HAVE_TFLITE
-
-static const TensorSpec Output =
-    TensorSpec::createSpec<float>(DecisionName, {1});
 static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
 
 #define _DECL_TRAIN_FEATURES(type, name, shape, _)                             \
   TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
 
 static const std::vector<TensorSpec> TrainingInputFeatures{
     {RA_PRIORITY_FEATURES_LIST(_DECL_TRAIN_FEATURES)
          TensorSpec::createSpec<float>("action_discount", {1}),
if (EC) {
       return false;
     }
     std::vector<TensorSpec> LFS = InputFeatures;
     if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
       append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
     // We always log the output; in particular, if we're not evaluating, we
     // don't have an output spec json file. That's why we handle the
     // 'normal' output separately.
-    LFS.push_back(Output);
+    LFS.push_back(DecisionSpec);
 
     Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
                                    /*IncludeReward*/ true);
     return false;
   }
 
   std::unique_ptr<RegAllocPriorityAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
private:
   std::unique_ptr<MLModelRunner> Runner;
   std::unique_ptr<Logger> Log;
 };
 #endif //#ifdef LLVM_HAVE_TFLITE
 
 } // namespace llvm
 
 RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() {
-  return new ReleaseModePriorityAdvisorAnalysis();
+  return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
+                 !InteractiveChannelBaseName.empty()
+             ? new ReleaseModePriorityAdvisorAnalysis()
+             : nullptr;
 }
 
 MLPriorityAdvisor::MLPriorityAdvisor(const MachineFunction &MF,
                                      const RAGreedy &RA,
                                      SlotIndexes *const Indexes,
                                      MLModelRunner *Runner)
     : RegAllocPriorityAdvisor(MF, RA, Indexes), DefaultAdvisor(MF, RA, Indexes),
       Runner(std::move(Runner)) {
   assert(this->Runner);
+  Runner->switchContext(MF.getName());
 }
 
 float MLPriorityAdvisor::getPriorityImpl(const LiveInterval &LI) const {
   const unsigned Size = LI.getSize();
   LiveRangeStage Stage = RA.getExtraInfo().getStage(LI);
 
   *Runner->getTensor<int64_t>(0) = static_cast<int64_t>(Size);
   *Runner->getTensor<int64_t>(1) = static_cast<int64_t>(Stage);

llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp

case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default:
     Ret = new DefaultEvictionAdvisorAnalysis(/*NotAsRequested*/ false);
     break;
   case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Development:
 #if defined(LLVM_HAVE_TFLITE)
     Ret = createDevelopmentModeAdvisor();
 #endif
     break;
   case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Release:
-#if defined(LLVM_HAVE_TF_AOT)
     Ret = createReleaseModeAdvisor();
-#endif
     break;
   }
   if (Ret)
     return Ret;
   return new DefaultEvictionAdvisorAnalysis(/*NotAsRequested*/ true);
 }
 
 StringRef RegAllocEvictionAdvisorAnalysis::getPassName() const {

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default:
     Ret = new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ false);
     break;
   case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development:
 #if defined(LLVM_HAVE_TFLITE)
     Ret = createDevelopmentModePriorityAdvisor();
 #endif
     break;
   case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release:
-#if defined(LLVM_HAVE_TF_AOT_REGALLOCPRIORITYMODEL)
     Ret = createReleaseModePriorityAdvisor();
-#endif
     break;
   }
   if (Ret)
     return Ret;
   return new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ true);
 }
 
 StringRef RegAllocPriorityAdvisorAnalysis::getPassName() const {

llvm/test/CodeGen/MLRegalloc/Inputs/interactive_main.py

+import log_reader
+import interactive_host
+import sys
+
+
+def main(args):
+  # this advisor just picks the first legal register to evict, which is
+  # identifiable by the "mask" feature
+  class Advisor:
+    to_return = False
+
+    def advice(self, tensor_values: list[log_reader.TensorValue]):
+      for tv in tensor_values:
+        if tv.spec().name != 'mask':
+          continue
+        for i, v in enumerate(tv):
+          if v == 1:
+            return i
+      # i.e. invalid:
+      return -1
+
+
+  a = Advisor()
+  interactive_host.run_interactive(args[0], a.advice, args[1:])
+
+
+if __name__ == '__main__':
+  main(sys.argv[1:])

llvm/test/CodeGen/MLRegalloc/interactive-mode.ll

+; 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:    llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=release -interactive-model-runner-echo-reply \
+; RUN:    -regalloc-evict-interactive-channel-base=%t.channel-basename %S/Inputs/two-large-fcts.ll -o /dev/null 2>%t.err | FileCheck %s
+; RUN: cat %t.err | FileCheck %s --check-prefix=ADVICE
+
+;; Make sure we see both contexts. Also sanity-check that the advice is the
+;; expected one - the index of the first legal register
+; CHECK: context: SyFgets
+; CHECK-NEXT: observation: 0
+; CHECK-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
+; CHECK: observation: 1
+; CHECK-NEXT: mask: 0,0,0,0,0,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
+; CHECK: context: SyFgetsCopy
+; CHECK-NEXT: observation: 0
+
+; ADVICE: index_to_evict: 9
+; ADVICE: index_to_evict: 10

llvm/test/Transforms/Inline/ML/Inputs/interactive_main.py

+import interactive_host
+import sys
+
+
+def main(args):
+
+  class Advisor:
+    to_return = False
+
+    def advice(self, _):
+      # The adice will be a sequence of yes/no/yes/no/...
+      # see ../interactive-mode.ll
+      self.to_return = not self.to_return
+      return int(self.to_return)
+
+  a = Advisor()
+  interactive_host.run_interactive(args[0], a.advice, args[1:])
+
+
+if __name__ == '__main__':
+  main(sys.argv[1:])

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

+; 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 2>%t.err | FileCheck %s
+; RUN: cat %t.err | FileCheck %s --check-prefix=ADVICE
+
+;; 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:      observation: 5
+; CHECK-NOT:  observation: 6
+
+; ADVICE:     inlining_decision: 1
+; ADVICE:     inlining_decision: 0
+; ADVICE:     inlining_decision: 1
+; ADVICE:     inlining_decision: 0
+; ADVICE:     inlining_decision: 1
+; ADVICE:     inlining_decision: 0