1//===- ReleaseModeModelRunner.h - 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#ifndef LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
15#define LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
16
17#include "llvm/Analysis/MLModelRunner.h"
18#include "llvm/Analysis/TensorSpec.h"
19#include "llvm/Support/ErrorHandling.h"
20
21#include <memory>
22#include <vector>
23
24namespace llvm {
25
26/// ReleaseModeModelRunner - production mode implementation of the
27/// MLModelRunner. It uses an AOT-compiled SavedModel for efficient execution.
28template <class TGen>
29class ReleaseModeModelRunner final : public MLModelRunner {
30public:
31 /// FeatureNames' type should be an indexed collection of std::string, like
32 /// std::array or std::vector, that has a size() method.
33 template <class FType>
34 ReleaseModeModelRunner(LLVMContext &Ctx, const FType &InputSpec,
35 StringRef DecisionName, StringRef FeedPrefix = "feed_",
36 StringRef FetchPrefix = "fetch_")
37 : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()),
38 CompiledModel(std::make_unique<TGen>()) {
39 assert(CompiledModel && "The CompiledModel should be valid");
40
41 for (size_t I = 0; I < InputSpec.size(); ++I) {
42 const int Index =
43 CompiledModel->LookupArgIndex(FeedPrefix.str() + InputSpec[I].name());
44 void *Buffer = nullptr;
45 if (Index >= 0)
46 Buffer = CompiledModel->arg_data(Index);
47 setUpBufferForTensor(Index: I, Spec: InputSpec[I], Buffer);
48 }
49
50 ResultIndex = CompiledModel->LookupResultIndex(FetchPrefix.str() +
51 DecisionName.str());
52 assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
53 }
54
55 virtual ~ReleaseModeModelRunner() = default;
56
57 static bool classof(const MLModelRunner *R) {
58 return R->getKind() == MLModelRunner::Kind::Release;
59 }
60
61private:
62 void *evaluateUntyped() override {
63 CompiledModel->Run();
64 return CompiledModel->result_data(ResultIndex);
65 }
66
67 int32_t ResultIndex = -1;
68 std::unique_ptr<TGen> CompiledModel;
69};
70
71/// A mock class satisfying the interface expected by ReleaseModeModelRunner for
72/// its `TGen` parameter. Useful to avoid conditional compilation complexity, as
73/// a compile-time replacement for a real AOT-ed model.
74class NoopSavedModelImpl final {
75#define NOOP_MODEL_ERRMSG \
76 "The mock AOT-ed saved model is a compile-time stub and should not be " \
77 "called."
78
79public:
80 NoopSavedModelImpl() = default;
81 int LookupArgIndex(const std::string &) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
82 int LookupResultIndex(const std::string &) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
83 void Run() { llvm_unreachable(NOOP_MODEL_ERRMSG); }
84 void *result_data(int) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
85 void *arg_data(int) { llvm_unreachable(NOOP_MODEL_ERRMSG); }
86#undef NOOP_MODEL_ERRMSG
87};
88
89template <class T> bool isEmbeddedModelEvaluatorValid() { return true; }
90
91template <> inline bool isEmbeddedModelEvaluatorValid<NoopSavedModelImpl>() {
92 return false;
93}
94} // namespace llvm
95
96#endif // LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
97

source code of llvm/include/llvm/Analysis/ReleaseModeModelRunner.h