diff --git a/llvm/include/llvm/Analysis/MLModelRunner.h b/llvm/include/llvm/Analysis/MLModelRunner.h --- a/llvm/include/llvm/Analysis/MLModelRunner.h +++ b/llvm/include/llvm/Analysis/MLModelRunner.h @@ -17,6 +17,9 @@ /// MLModelRunner interface: abstraction of a mechanism for evaluating a /// tensorflow "saved model". +/// NOTE: feature indices are expected to be consistent all accross +/// MLModelRunners (pertaining to the same model), and also Loggers (see +/// TFUtils.h) class MLModelRunner { public: // Disallows copy and assign. diff --git a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h @@ -0,0 +1,67 @@ +//===- ReleaseModeModelRunner.h - Fast, precompiled model runner ---------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file implements a model runner wrapping an AOT compiled ML model. +// Only inference is supported. +// +//===----------------------------------------------------------------------===// +#include "llvm/Analysis/MLModelRunner.h" + +#include +#include + +using namespace llvm; +namespace llvm { + +/// ReleaseModeModelRunner - production mode implementation of the +/// MLModelRunner. It uses an AOT-compiled SavedModel for efficient execution. +template +class ReleaseModeModelRunner final : public MLModelRunner { +public: + /// FeatureNames' type should be an indexed collection of std::string, like + /// std::array or std::vector, that has a size() method. + template + ReleaseModeModelRunner(LLVMContext &Ctx, const FType &FeatureNames, + StringRef DecisionName, StringRef FeedPrefix = "feed_", + StringRef FetchPrefix = "fetch_") + : MLModelRunner(Ctx), CompiledModel(std::make_unique()) { + assert(CompiledModel && "The CompiledModel should be valid"); + + const size_t FeatureCount = FeatureNames.size(); + FeatureIndices.resize(FeatureCount); + + for (size_t I = 0; I < FeatureCount; ++I) { + const int Index = + CompiledModel->LookupArgIndex(FeedPrefix.str() + FeatureNames[I]); + assert(Index >= 0 && "Cannot find Feature in inlining model"); + FeatureIndices[I] = Index; + } + + ResultIndex = CompiledModel->LookupResultIndex(FetchPrefix.str() + + DecisionName.str()); + assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model"); + } + + virtual ~ReleaseModeModelRunner() = default; + +private: + void *evaluateUntyped() override { + CompiledModel->Run(); + return CompiledModel->result_data(ResultIndex); + } + + void *getTensorUntyped(size_t Index) override { + return reinterpret_cast( + CompiledModel->arg_data(FeatureIndices[Index])); + } + + std::vector FeatureIndices; + int32_t ResultIndex = -1; + std::unique_ptr CompiledModel; +}; +} // namespace llvm diff --git a/llvm/include/llvm/Analysis/Utils/TFUtils.h b/llvm/include/llvm/Analysis/Utils/TFUtils.h --- a/llvm/include/llvm/Analysis/Utils/TFUtils.h +++ b/llvm/include/llvm/Analysis/Utils/TFUtils.h @@ -147,6 +147,9 @@ /// Construct a Logger. If IncludeReward is false, then logReward or /// logFinalReward shouldn't be called, and the reward feature won't be /// printed out. + /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have + /// corresponding indices) with any MLModelRunner implementations + /// corresponding to the model being trained/logged. Logger(const std::vector &FeatureSpecs, const TensorSpec &RewardSpec, bool IncludeReward); diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt --- a/llvm/lib/Analysis/CMakeLists.txt +++ b/llvm/lib/Analysis/CMakeLists.txt @@ -121,7 +121,6 @@ RegionInfo.cpp RegionPass.cpp RegionPrinter.cpp - ReleaseModeModelRunner.cpp ReplayInlineAdvisor.cpp ScalarEvolution.cpp ScalarEvolutionAliasAnalysis.cpp diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp --- a/llvm/lib/Analysis/MLInlineAdvisor.cpp +++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp @@ -35,6 +35,21 @@ using namespace llvm; +#ifdef LLVM_HAVE_TF_AOT +#include "llvm/Analysis/ReleaseModeModelRunner.h" +// codegen-ed file +#include "InlinerSizeModel.h" // NOLINT +#include "llvm/Analysis/InlineModelFeatureMaps.h" + +std::unique_ptr +llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) { + auto AOTRunner = + std::make_unique>( + M.getContext(), FeatureNameMap, DecisionName); + return std::make_unique(M, MAM, std::move(AOTRunner)); +} +#endif + #define DEBUG_TYPE "inline-ml" static cl::opt SizeIncreaseThreshold( diff --git a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp deleted file mode 100644 --- a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp +++ /dev/null @@ -1,82 +0,0 @@ -//===- ReleaseModeModelRunner.cpp - Fast, precompiled model runner -------===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// -// -// This file implements a model runner wrapping an AOT compiled ML model. -// Only inference is supported. -// -//===----------------------------------------------------------------------===// -#include "llvm/Config/config.h" -#if defined(LLVM_HAVE_TF_AOT) - -#include "llvm/Analysis/InlineModelFeatureMaps.h" -#include "llvm/Analysis/MLInlineAdvisor.h" - -// codegen-ed file -#include "InlinerSizeModel.h" // NOLINT - -#include -#include - -using namespace llvm; -namespace { - -const char FeedPrefix[] = "feed_"; -const char FetchPrefix[] = "fetch_"; - -/// MLModelRunner - production mode implementation. It uses a AOT-compiled -/// SavedModel for efficient execution. -class ReleaseModeModelRunner final : public MLModelRunner { -public: - ReleaseModeModelRunner(LLVMContext &Ctx); - virtual ~ReleaseModeModelRunner() = default; - -private: - void *evaluateUntyped() override; - void *getTensorUntyped(size_t Index) override; - - std::vector FeatureIndices; - int32_t ResultIndex = -1; - std::unique_ptr CompiledModel; -}; -} // namespace - -ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx) - : MLModelRunner(Ctx), - CompiledModel(std::make_unique()) { - assert(CompiledModel && "The CompiledModel should be valid"); - - FeatureIndices.resize(NumberOfFeatures); - - for (size_t I = 0; I < NumberOfFeatures; ++I) { - const int Index = - CompiledModel->LookupArgIndex(FeedPrefix + FeatureNameMap[I]); - assert(Index >= 0 && "Cannot find Feature in inlining model"); - FeatureIndices[I] = Index; - } - - ResultIndex = - CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName); - assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model"); -} - -void *ReleaseModeModelRunner::getTensorUntyped(size_t Index) { - return reinterpret_cast( - CompiledModel->arg_data(FeatureIndices[Index])); -} - -void *ReleaseModeModelRunner::evaluateUntyped() { - CompiledModel->Run(); - return CompiledModel->result_data(ResultIndex); -} - -std::unique_ptr -llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) { - auto AOTRunner = std::make_unique(M.getContext()); - return std::make_unique(M, MAM, std::move(AOTRunner)); -} -#endif // defined(LLVM_HAVE_TF_AOT)