1//===- ReleaseModeModelRunner.cpp - Fast, precompiled model runner  -------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements a model runner wrapping an AOT compiled ML model.
10// Only inference is supported.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/Analysis/InlineModelFeatureMaps.h"
15#include "llvm/Analysis/MLInlineAdvisor.h"
16
17// codegen-ed file
18#include "InlinerSizeModel.h" // NOLINT
19
20#include <memory>
21#include <vector>
22
23using namespace llvm;
24namespace {
25
26static const char *const FeedPrefix = "feed_";
27static const char *const FetchPrefix = "fetch_";
28
29/// MLModelRunner - production mode implementation. It uses a AOT-compiled
30/// SavedModel for efficient execution.
31class ReleaseModeModelRunner final : public MLModelRunner {
32public:
33  ReleaseModeModelRunner(LLVMContext &Ctx);
34  virtual ~ReleaseModeModelRunner() = default;
35
36  bool run() override;
37
38  void setFeature(FeatureIndex Index, int64_t Value) override;
39  int64_t getFeature(int Index) const override;
40
41private:
42  std::vector<int32_t> FeatureIndices;
43  int32_t ResultIndex = -1;
44  std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
45};
46} // namespace
47
48ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
49    : MLModelRunner(Ctx),
50      CompiledModel(std::make_unique<llvm::InlinerSizeModel>()) {
51  assert(CompiledModel && "The CompiledModel should be valid");
52
53  FeatureIndices.reserve(NumberOfFeatures);
54
55  for (size_t I = 0; I < NumberOfFeatures; ++I) {
56    const int Index =
57        CompiledModel->LookupArgIndex(FeedPrefix + FeatureNameMap[I]);
58    assert(Index >= 0 && "Cannot find Feature in inlining model");
59    FeatureIndices[I] = Index;
60  }
61
62  ResultIndex =
63      CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName);
64  assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
65}
66
67int64_t ReleaseModeModelRunner::getFeature(int Index) const {
68  return *static_cast<int64_t *>(
69      CompiledModel->arg_data(FeatureIndices[Index]));
70}
71
72void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
73  *static_cast<int64_t *>(CompiledModel->arg_data(
74      FeatureIndices[static_cast<size_t>(Index)])) = Value;
75}
76
77bool ReleaseModeModelRunner::run() {
78  CompiledModel->Run();
79  return static_cast<bool>(
80      *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
81}
82
83std::unique_ptr<InlineAdvisor>
84llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
85  auto AOTRunner = std::make_unique<ReleaseModeModelRunner>(M.getContext());
86  return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
87}
88