[llvm] db5aceb - [NFC] Expose the ReleaseModeModelRunner
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Wed Dec 15 23:22:07 PST 2021
Author: Mircea Trofin
Date: 2021-12-15T23:21:58-08:00
New Revision: db5aceb9793967992c1bcf154327979a4328ec42
URL: https://github.com/llvm/llvm-project/commit/db5aceb9793967992c1bcf154327979a4328ec42
DIFF: https://github.com/llvm/llvm-project/commit/db5aceb9793967992c1bcf154327979a4328ec42.diff
LOG: [NFC] Expose the ReleaseModeModelRunner
The type was pretty much generic, just needed a bit of parameterization.
Differential Revision: https://reviews.llvm.org/D115764
Added:
llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
Modified:
llvm/include/llvm/Analysis/MLModelRunner.h
llvm/include/llvm/Analysis/Utils/TFUtils.h
llvm/lib/Analysis/CMakeLists.txt
llvm/lib/Analysis/MLInlineAdvisor.cpp
Removed:
llvm/lib/Analysis/ReleaseModeModelRunner.cpp
################################################################################
diff --git a/llvm/include/llvm/Analysis/MLModelRunner.h b/llvm/include/llvm/Analysis/MLModelRunner.h
index 2281b63448354..90b3cc7e76e65 100644
--- a/llvm/include/llvm/Analysis/MLModelRunner.h
+++ b/llvm/include/llvm/Analysis/MLModelRunner.h
@@ -17,6 +17,9 @@ namespace llvm {
/// 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
index 0000000000000..b684f87ea5cba
--- /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 <memory>
+#include <vector>
+
+using namespace llvm;
+namespace llvm {
+
+/// ReleaseModeModelRunner - production mode implementation of the
+/// MLModelRunner. It uses an AOT-compiled SavedModel for efficient execution.
+template <class TGen>
+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 <class FType>
+ ReleaseModeModelRunner(LLVMContext &Ctx, const FType &FeatureNames,
+ StringRef DecisionName, StringRef FeedPrefix = "feed_",
+ StringRef FetchPrefix = "fetch_")
+ : MLModelRunner(Ctx), CompiledModel(std::make_unique<TGen>()) {
+ 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<char *>(
+ CompiledModel->arg_data(FeatureIndices[Index]));
+ }
+
+ std::vector<int32_t> FeatureIndices;
+ int32_t ResultIndex = -1;
+ std::unique_ptr<TGen> CompiledModel;
+};
+} // namespace llvm
diff --git a/llvm/include/llvm/Analysis/Utils/TFUtils.h b/llvm/include/llvm/Analysis/Utils/TFUtils.h
index e4a4b7428cb6c..012fca53a2002 100644
--- a/llvm/include/llvm/Analysis/Utils/TFUtils.h
+++ b/llvm/include/llvm/Analysis/Utils/TFUtils.h
@@ -147,6 +147,9 @@ class Logger final {
/// 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<LoggedFeatureSpec> &FeatureSpecs,
const TensorSpec &RewardSpec, bool IncludeReward);
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index ba5735061897b..ae0ee22561ced 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -121,7 +121,6 @@ add_llvm_component_library(LLVMAnalysis
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
index 17f6fa5034567..f5a65cd2b689d 100644
--- 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<InlineAdvisor>
+llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
+ auto AOTRunner =
+ std::make_unique<ReleaseModeModelRunner<llvm::InlinerSizeModel>>(
+ M.getContext(), FeatureNameMap, DecisionName);
+ return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
+}
+#endif
+
#define DEBUG_TYPE "inline-ml"
static cl::opt<float> SizeIncreaseThreshold(
diff --git a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp
deleted file mode 100644
index 3d4c75dd4d843..0000000000000
--- 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 <memory>
-#include <vector>
-
-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<int32_t> FeatureIndices;
- int32_t ResultIndex = -1;
- std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
-};
-} // namespace
-
-ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
- : MLModelRunner(Ctx),
- CompiledModel(std::make_unique<llvm::InlinerSizeModel>()) {
- 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<char *>(
- CompiledModel->arg_data(FeatureIndices[Index]));
-}
-
-void *ReleaseModeModelRunner::evaluateUntyped() {
- CompiledModel->Run();
- return CompiledModel->result_data(ResultIndex);
-}
-
-std::unique_ptr<InlineAdvisor>
-llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
- auto AOTRunner = std::make_unique<ReleaseModeModelRunner>(M.getContext());
- return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
-}
-#endif // defined(LLVM_HAVE_TF_AOT)
More information about the llvm-commits
mailing list