859ebca744
This reverts commit 640beb38e7710b939b3cfb3f4c54accc694b1d30. That commit caused performance degradtion in Quicksilver test QS:sGPU and a functional test failure in (rocPRIM rocprim.device_segmented_radix_sort). Reverting until we have a better solution to s_cselect_b64 codegen cleanup Change-Id: Ibf8e397df94001f248fba609f072088a46abae08 Reviewed By: kzhuravl Differential Revision: https://reviews.llvm.org/D115960 Change-Id: Id169459ce4dfffa857d5645a0af50b0063ce1105
91 lines
2.9 KiB
C++
91 lines
2.9 KiB
C++
//===- 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;
|
|
|
|
bool run() override;
|
|
|
|
void setFeature(FeatureIndex Index, int64_t Value) override;
|
|
int64_t getFeature(int Index) const override;
|
|
|
|
private:
|
|
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");
|
|
}
|
|
|
|
int64_t ReleaseModeModelRunner::getFeature(int Index) const {
|
|
return *static_cast<int64_t *>(
|
|
CompiledModel->arg_data(FeatureIndices[Index]));
|
|
}
|
|
|
|
void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
|
|
*static_cast<int64_t *>(CompiledModel->arg_data(
|
|
FeatureIndices[static_cast<size_t>(Index)])) = Value;
|
|
}
|
|
|
|
bool ReleaseModeModelRunner::run() {
|
|
CompiledModel->Run();
|
|
return static_cast<bool>(
|
|
*static_cast<int64_t *>(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)
|