Details | Last modification | View Log | RSS feed
| Rev | Author | Line No. | Line |
|---|---|---|---|
| 14 | pmbaty | 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(I, 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 | } // namespace llvm |
||
| 89 | |||
| 90 | #endif // LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H |