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

This reverts commit a772f0b.

The main problem was related to how we handled `dbgs()` from the hosted
compiler. Using explicit `subprocess.communicate`, and not relying on
dbgs() being flushed until the end appears to address the problem.

Also some fixes due to some bots running older pythons, so we can't have
nice things like `int | float` and such.

Author: mtrofin

llvm/include/llvm/Analysis/InlineModelFeatureMaps.h

@@ -129,9 +129,10 @@ inlineCostFeatureToMlFeature(InlineCostFeatureIndex 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;
 

llvm/include/llvm/Analysis/InteractiveModelRunner.h

@@ -48,7 +48,7 @@ class InteractiveModelRunner : public MLModelRunner {
   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();
   }

llvm/include/llvm/Analysis/MLModelRunner.h

@@ -49,6 +49,7 @@ class MLModelRunner {
 
   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)

llvm/include/llvm/Analysis/ReleaseModeModelRunner.h

@@ -85,6 +85,12 @@ class NoopSavedModelImpl final {
   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

@@ -283,7 +283,7 @@ TrainingLogger::TrainingLogger(StringRef LogFileName,
   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)

llvm/lib/Analysis/InlineAdvisor.cpp

@@ -231,10 +231,8 @@ bool InlineAdvisorAnalysis::Result::tryCreate(
 #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;
   }
 

llvm/lib/Analysis/MLInlineAdvisor.cpp

@@ -18,10 +18,12 @@
 #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"
@@ -30,19 +32,37 @@
 
 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>>(
+        M.getContext(), FeatureMap, DecisionName);
+  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"
 
@@ -59,7 +79,7 @@ static cl::opt<bool> KeepFPICache(
     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)
@@ -73,6 +93,8 @@ const std::array<TensorSpec, NumberOfFeatures> llvm::FeatureMap{
 // 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";
 
@@ -94,7 +116,7 @@ MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
       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

llvm/lib/Analysis/models/interactive_host.py

@@ -0,0 +1,84 @@
+"""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, Union
+
+
+def send(f: io.BufferedWriter, value: Union[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]],
+                                            Union[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, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL)
+    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)
+    _, err = compiler_proc.communicate()
+    print(err.decode('utf-8'))
+    compiler_proc.wait()
+
+  finally:
+    os.unlink(to_compiler)
+    os.unlink(from_compiler)

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

@@ -13,6 +13,7 @@
 #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)
@@ -52,6 +53,14 @@ using CompiledModelType = RegallocEvictModel;
 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"
@@ -213,6 +222,8 @@ static const std::vector<int64_t> PerLiveRangeShape{1, NumberOfInterferences};
 // 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 {
@@ -382,14 +393,21 @@ class ReleaseModeEvictionAdvisorAnalysis final
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-    if (!Runner)
-      Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
-          MF.getFunction().getContext(), InputFeatures, DecisionName);
+    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;
 };
 
 // ===================================
@@ -398,8 +416,6 @@ class ReleaseModeEvictionAdvisorAnalysis final
 //
 // 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
@@ -512,7 +528,7 @@ class DevelopmentModeEvictionAdvisorAnalysis final
     // 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);
@@ -557,6 +573,7 @@ MLEvictAdvisor::MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &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);
@@ -1134,7 +1151,10 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
 #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.