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