1 files changed, 87 insertions, 0 deletions
diff --git a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp
new file mode 100644
index 0000000000000..4c0ffbc17ff7a
--- /dev/null
+++ b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp
@@ -0,0 +1,87 @@
+//===- ReleaseModeModelRunner.cpp - Fast, precompiled model runner  -------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a model runner wrapping an AOT compiled ML model.
+// Only inference is supported.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/InlineModelFeatureMaps.h"
+#include "llvm/Analysis/MLInlineAdvisor.h"
+
+// codegen-ed file
+#include "InlinerSizeModel.h" // NOLINT
+
+#include <memory>
+#include <vector>
+
+using namespace llvm;
+namespace {
+
+static const char *const FeedPrefix = "feed_";
+static const char *const FetchPrefix = "fetch_";
+
+/// MLModelRunner - production mode implementation. It uses a AOT-compiled
+/// SavedModel for efficient execution.
+class ReleaseModeModelRunner final : public MLModelRunner {
+public:
+  ReleaseModeModelRunner(LLVMContext &Ctx);
+  virtual ~ReleaseModeModelRunner() = default;
+
+  bool run() override;
+
+  void setFeature(FeatureIndex Index, int64_t Value) override;
+  int64_t getFeature(int Index) const override;
+
+private:
+  std::vector<int32_t> FeatureIndices;
+  int32_t ResultIndex = -1;
+  std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
+};
+} // namespace
+
+ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
+    : MLModelRunner(Ctx),
+      CompiledModel(std::make_unique<llvm::InlinerSizeModel>()) {
+  assert(CompiledModel && "The CompiledModel should be valid");
+
+  FeatureIndices.reserve(NumberOfFeatures);
+
+  for (size_t I = 0; I < NumberOfFeatures; ++I) {
+    const int Index =
+        CompiledModel->LookupArgIndex(FeedPrefix + FeatureNameMap[I]);
+    assert(Index >= 0 && "Cannot find Feature in inlining model");
+    FeatureIndices[I] = Index;
+  }
+
+  ResultIndex =
+      CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName);
+  assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
+}
+
+int64_t ReleaseModeModelRunner::getFeature(int Index) const {
+  return *static_cast<int64_t *>(
+      CompiledModel->arg_data(FeatureIndices[Index]));
+}
+
+void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
+  *static_cast<int64_t *>(CompiledModel->arg_data(
+      FeatureIndices[static_cast<size_t>(Index)])) = Value;
+}
+
+bool ReleaseModeModelRunner::run() {
+  CompiledModel->Run();
+  return static_cast<bool>(
+      *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
+}
+
+std::unique_ptr<InlineAdvisor>
+llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
+  auto AOTRunner = std::make_unique<ReleaseModeModelRunner>(M.getContext());
+  return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
+}