llvm/lib/Analysis/ReleaseModeModelRunner.cpp

   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.resize(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)