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 #include "llvm/Config/config.h"
14 #if defined(LLVM_HAVE_TF_AOT)
15
16 #include "llvm/Analysis/InlineModelFeatureMaps.h"
17 #include "llvm/Analysis/MLInlineAdvisor.h"
18
19 // codegen-ed file
20 #include "InlinerSizeModel.h" // NOLINT
21
22 #include <memory>
23 #include <vector>
24
25 using namespace llvm;
26 namespace {
27
28 const char FeedPrefix[] = "feed_";
29 const char FetchPrefix[] = "fetch_";
30
31 /// MLModelRunner - production mode implementation. It uses a AOT-compiled
32 /// SavedModel for efficient execution.
33 class ReleaseModeModelRunner final : public MLModelRunner {
34 public:
35 ReleaseModeModelRunner(LLVMContext &Ctx);
36 virtual ~ReleaseModeModelRunner() = default;
37
38 bool run() override;
39
40 void setFeature(FeatureIndex Index, int64_t Value) override;
41 int64_t getFeature(int Index) const override;
42
43 private:
44 std::vector<int32_t> FeatureIndices;
45 int32_t ResultIndex = -1;
46 std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
47 };
48 } // namespace
49
ReleaseModeModelRunner(LLVMContext & Ctx)50 ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
51 : MLModelRunner(Ctx),
52 CompiledModel(std::make_unique<llvm::InlinerSizeModel>()) {
53 assert(CompiledModel && "The CompiledModel should be valid");
54
55 FeatureIndices.reserve(NumberOfFeatures);
56
57 for (size_t I = 0; I < NumberOfFeatures; ++I) {
58 const int Index =
59 CompiledModel->LookupArgIndex(FeedPrefix + FeatureNameMap[I]);
60 assert(Index >= 0 && "Cannot find Feature in inlining model");
61 FeatureIndices[I] = Index;
62 }
63
64 ResultIndex =
65 CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName);
66 assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
67 }
68
getFeature(int Index) const69 int64_t ReleaseModeModelRunner::getFeature(int Index) const {
70 return *static_cast<int64_t *>(
71 CompiledModel->arg_data(FeatureIndices[Index]));
72 }
73
setFeature(FeatureIndex Index,int64_t Value)74 void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
75 *static_cast<int64_t *>(CompiledModel->arg_data(
76 FeatureIndices[static_cast<size_t>(Index)])) = Value;
77 }
78
run()79 bool ReleaseModeModelRunner::run() {
80 CompiledModel->Run();
81 return static_cast<bool>(
82 *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
83 }
84
85 std::unique_ptr<InlineAdvisor>
getReleaseModeAdvisor(Module & M,ModuleAnalysisManager & MAM)86 llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
87 auto AOTRunner = std::make_unique<ReleaseModeModelRunner>(M.getContext());
88 return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
89 }
90 #endif // defined(LLVM_HAVE_TF_AOT)
91