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 | |
24 | namespace llvm { |
25 | |
26 | /// ReleaseModeModelRunner - production mode implementation of the |
27 | /// MLModelRunner. It uses an AOT-compiled SavedModel for efficient execution. |
28 | template <class TGen> |
29 | class ReleaseModeModelRunner final : public MLModelRunner { |
30 | public: |
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 | |
61 | private: |
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. |
74 | class 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 | |
79 | public: |
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 | |
89 | template <class T> bool isEmbeddedModelEvaluatorValid() { return true; } |
90 | |
91 | template <> inline bool isEmbeddedModelEvaluatorValid<NoopSavedModelImpl>() { |
92 | return false; |
93 | } |
94 | } // namespace llvm |
95 | |
96 | #endif // LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H |
97 | |