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 
23 using namespace llvm;
24 namespace {
25 
26 static const char *const FeedPrefix = "feed_";
27 static const char *const FetchPrefix = "fetch_";
28 
29 /// MLModelRunner - production mode implementation. It uses a AOT-compiled
30 /// SavedModel for efficient execution.
31 class ReleaseModeModelRunner final : public MLModelRunner {
32 public:
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 
41 private:
42   std::vector<int32_t> FeatureIndices;
43   int32_t ResultIndex = -1;
44   std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
45 };
46 } // namespace
47 
48 ReleaseModeModelRunner::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 
67 int64_t ReleaseModeModelRunner::getFeature(int Index) const {
68   return *static_cast<int64_t *>(
69       CompiledModel->arg_data(FeatureIndices[Index]));
70 }
71 
72 void 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 
77 bool ReleaseModeModelRunner::run() {
78   CompiledModel->Run();
79   return static_cast<bool>(
80       *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
81 }
82 
83 std::unique_ptr<InlineAdvisor>
84 llvm::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