Home | History | Annotate | Line # | Download | only in Analysis 
      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 
     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 
     69 int64_t ReleaseModeModelRunner::getFeature(int Index) const {
     70   return *static_cast<int64_t *>(
     71       CompiledModel->arg_data(FeatureIndices[Index]));
     72 }
     73 
     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 
     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>
     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