From 059e03476cbbb7349b75fa9cc26bc9da0d491529 Mon Sep 17 00:00:00 2001 From: Mircea Trofin Date: Tue, 7 Dec 2021 15:07:39 -0800 Subject: [PATCH] [NFC][mlgo] Generalize model runner interface This prepares it for the regalloc work. Part of it is making model evaluation accross 'development' and 'release' scenarios more reusable. This patch: - extends support to tensors of any shape (not just scalars, like we had in the inliner -Oz case). While the tensor shape can be anything, we assume row-major layout and expose the tensor as a buffer. - exposes the NoInferenceModelRunner, which we use in the 'development' mode to keep the evaluation code path consistent and simplify logging, as we'll want to reuse it in the regalloc case. Differential Revision: https://reviews.llvm.org/D115306 --- llvm/include/llvm/Analysis/MLModelRunner.h | 22 +++++- .../llvm/Analysis/NoInferenceModelRunner.h | 39 +++++++++ llvm/include/llvm/Analysis/Utils/TFUtils.h | 4 +- llvm/lib/Analysis/CMakeLists.txt | 1 + .../Analysis/DevelopmentModeInlineAdvisor.cpp | 79 ++++++++----------- llvm/lib/Analysis/MLInlineAdvisor.cpp | 47 ++++++----- llvm/lib/Analysis/NoInferenceModelRunner.cpp | 33 ++++++++ llvm/lib/Analysis/ReleaseModeModelRunner.cpp | 22 ++---- llvm/unittests/Analysis/CMakeLists.txt | 5 +- llvm/unittests/Analysis/MLModelRunnerTest.cpp | 33 ++++++++ 10 files changed, 195 insertions(+), 90 deletions(-) create mode 100644 llvm/include/llvm/Analysis/NoInferenceModelRunner.h create mode 100644 llvm/lib/Analysis/NoInferenceModelRunner.cpp create mode 100644 llvm/unittests/Analysis/MLModelRunnerTest.cpp diff --git a/llvm/include/llvm/Analysis/MLModelRunner.h b/llvm/include/llvm/Analysis/MLModelRunner.h index 7cfa6efedf10..2281b6344835 100644 --- a/llvm/include/llvm/Analysis/MLModelRunner.h +++ b/llvm/include/llvm/Analysis/MLModelRunner.h @@ -10,7 +10,6 @@ #ifndef LLVM_ANALYSIS_MLMODELRUNNER_H #define LLVM_ANALYSIS_MLMODELRUNNER_H -#include "llvm/Analysis/InlineModelFeatureMaps.h" #include "llvm/IR/LLVMContext.h" #include "llvm/IR/PassManager.h" @@ -25,12 +24,27 @@ public: MLModelRunner &operator=(const MLModelRunner &) = delete; virtual ~MLModelRunner() = default; - virtual bool run() = 0; - virtual void setFeature(FeatureIndex Index, int64_t Value) = 0; - virtual int64_t getFeature(int Index) const = 0; + template T evaluate() { + return *reinterpret_cast(evaluateUntyped()); + } + + template T *getTensor(I FeatureID) { + return reinterpret_cast( + getTensorUntyped(static_cast(FeatureID))); + } + + template const T *getTensor(I FeatureID) const { + return reinterpret_cast( + getTensorUntyped(static_cast(FeatureID))); + } protected: MLModelRunner(LLVMContext &Ctx) : Ctx(Ctx) {} + virtual void *evaluateUntyped() = 0; + virtual void *getTensorUntyped(size_t Index) = 0; + const void *getTensorUntyped(size_t Index) const { + return (const_cast(this))->getTensorUntyped(Index); + } LLVMContext &Ctx; }; diff --git a/llvm/include/llvm/Analysis/NoInferenceModelRunner.h b/llvm/include/llvm/Analysis/NoInferenceModelRunner.h new file mode 100644 index 000000000000..6f6f5a009b1c --- /dev/null +++ b/llvm/include/llvm/Analysis/NoInferenceModelRunner.h @@ -0,0 +1,39 @@ +//===- NoInferenceModelRunner.h ---- noop ML model runner ------*- C++ -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// + +#ifndef LLVM_ANALYSIS_NOINFERENCEMODELRUNNER_H +#define LLVM_ANALYSIS_NOINFERENCEMODELRUNNER_H + +#include "llvm/Config/llvm-config.h" + +/// While not strictly necessary to conditionally compile this, it really +/// has no usecase outside the 'development' mode. +#ifdef LLVM_HAVE_TF_API +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/Utils/TFUtils.h" +namespace llvm { +/// A pseudo model runner. We use it to store feature values when collecting +/// logs for the default policy, in 'development' mode, but never ask it to +/// 'run'. +class NoInferenceModelRunner : public MLModelRunner { +public: + NoInferenceModelRunner(LLVMContext &Ctx, + const std::vector &Inputs); + +private: + void *evaluateUntyped() override { + llvm_unreachable("We shouldn't call run on this model runner."); + } + void *getTensorUntyped(size_t Index) override; + + std::vector> ValuesBuffer; +}; +} // namespace llvm +#endif // defined(LLVM_HAVE_TF_API) +#endif // defined(LLVM_ANALYSIS_NOINFERENCEMODELRUNNER_H) \ No newline at end of file diff --git a/llvm/include/llvm/Analysis/Utils/TFUtils.h b/llvm/include/llvm/Analysis/Utils/TFUtils.h index 1f6be0e60eb9..e4a4b7428cb6 100644 --- a/llvm/include/llvm/Analysis/Utils/TFUtils.h +++ b/llvm/include/llvm/Analysis/Utils/TFUtils.h @@ -246,8 +246,10 @@ public: /// otherwise. bool isValid() const { return !!Impl; } -private: + /// Untyped access to input. void *getUntypedInput(size_t Index); + +private: std::unique_ptr Impl; }; diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt index 9da07cb1c485..9dfc619f2b13 100644 --- a/llvm/lib/Analysis/CMakeLists.txt +++ b/llvm/lib/Analysis/CMakeLists.txt @@ -105,6 +105,7 @@ add_llvm_component_library(LLVMAnalysis ModuleDebugInfoPrinter.cpp ModuleSummaryAnalysis.cpp MustExecute.cpp + NoInferenceModelRunner.cpp ObjCARCAliasAnalysis.cpp ObjCARCAnalysisUtils.cpp ObjCARCInstKind.cpp diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp index d87fa849d839..b76d63dddeb0 100644 --- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp +++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp @@ -16,6 +16,7 @@ #include "llvm/Analysis/CallGraph.h" #include "llvm/Analysis/InlineSizeEstimatorAnalysis.h" #include "llvm/Analysis/MLInlineAdvisor.h" +#include "llvm/Analysis/NoInferenceModelRunner.h" #include "llvm/Analysis/Utils/TFUtils.h" #include "llvm/IR/LLVMContext.h" #include "llvm/Support/CommandLine.h" @@ -261,25 +262,6 @@ private: const int64_t Mandatory; }; -/// A pseudo model runner. We use it to store feature values when collecting -/// logs for the default policy, but never ask it to 'run'. -class NoInferenceModelRunner : public MLModelRunner { -public: - NoInferenceModelRunner(LLVMContext &Ctx) - : MLModelRunner(Ctx), Features(NumberOfFeatures) {} - void setFeature(FeatureIndex Index, int64_t Value) override { - Features[static_cast(Index)] = Value; - } - - int64_t getFeature(int Index) const override { return Features[Index]; } - bool run() override { - llvm_unreachable("We shouldn't call run on this model runner."); - } - -private: - InlineFeatures Features; -}; - /// ModelUnderTrainingRunner - training mode implementation. It uses TF C APIs /// to dynamically load and evaluate a TF SavedModel /// (https://www.tensorflow.org/guide/saved_model). Runtime performance is @@ -288,15 +270,11 @@ class ModelUnderTrainingRunner final : public MLModelRunner { public: ModelUnderTrainingRunner(LLVMContext &Ctx, const std::string &ModelPath); - bool run() override; - // Disallows copy and assign. ModelUnderTrainingRunner(const ModelUnderTrainingRunner &) = delete; ModelUnderTrainingRunner & operator=(const ModelUnderTrainingRunner &) = delete; - void setFeature(FeatureIndex Index, int64_t Value) override; - int64_t getFeature(int Index) const override; bool isValid() const { return !!Evaluator; } const std::vector &outputLoggedFeatureSpecs() const { @@ -308,18 +286,31 @@ public: return LastEvaluationResult; } + static const std::vector getInputFeatures() { + std::vector InputSpecs; + for (size_t I = 0; I < NumberOfFeatures; ++I) + InputSpecs.push_back(TensorSpec::createSpec( + TFFeedPrefix + FeatureNameMap[I], {1})); + append_range(InputSpecs, TrainingOnlyFeatures); + return InputSpecs; + } + private: std::unique_ptr Evaluator; std::vector OutputSpecs; Optional LastEvaluationResult; + void *evaluateUntyped() override; + void *getTensorUntyped(size_t Index) override; // The training framework needs some additional features. - const std::vector TrainingOnlyFeatures{ - TensorSpec::createSpec(TFFeedPrefix + "inlining_default", {1}), - TensorSpec::createSpec(TFFeedPrefix + "discount", {1}), - TensorSpec::createSpec(TFFeedPrefix + "reward", {1}), - TensorSpec::createSpec(TFFeedPrefix + "step_type", {1})}; + const static std::vector TrainingOnlyFeatures; }; + +const std::vector ModelUnderTrainingRunner::TrainingOnlyFeatures{ + TensorSpec::createSpec(TFFeedPrefix + "inlining_default", {1}), + TensorSpec::createSpec(TFFeedPrefix + "discount", {1}), + TensorSpec::createSpec(TFFeedPrefix + "reward", {1}), + TensorSpec::createSpec(TFFeedPrefix + "step_type", {1})}; } // namespace TrainingLogger::TrainingLogger(StringRef LogFileName, @@ -353,7 +344,7 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event, const MLModelRunner &ModelRunner) { size_t CurrentFeature = 0; for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) { - int64_t F = ModelRunner.getFeature(CurrentFeature); + int64_t F = *ModelRunner.getTensor(CurrentFeature); L->logInt64Value(CurrentFeature, &F); } @@ -433,7 +424,9 @@ DevelopmentModeMLInlineAdvisor::getAdviceFromModel( return MLInlineAdvisor::getAdviceFromModel(CB, ORE); bool DefaultAdvice = GetDefaultAdvice(CB); - auto Recommendation = IsDoingInference ? ModelRunner->run() : DefaultAdvice; + auto Recommendation = + IsDoingInference ? static_cast(ModelRunner->evaluate()) + : DefaultAdvice; return std::make_unique( /*Advisor=*/this, /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation, @@ -461,11 +454,8 @@ size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() { ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx, const std::string &ModelPath) : MLModelRunner(Ctx) { - std::vector InputSpecs; - for (size_t I = 0; I < NumberOfFeatures; ++I) - InputSpecs.push_back( - TensorSpec::createSpec(TFFeedPrefix + FeatureNameMap[I], {1})); - append_range(InputSpecs, TrainingOnlyFeatures); + std::vector InputSpecs = + ModelUnderTrainingRunner::getInputFeatures(); if (auto MaybeOutSpecs = loadOutputSpecs(Ctx, DecisionName, ModelPath, TFOutputSpecOverride)) OutputSpecs = std::move(*MaybeOutSpecs); @@ -482,23 +472,17 @@ ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx, } } -bool ModelUnderTrainingRunner::run() { +void *ModelUnderTrainingRunner::evaluateUntyped() { LastEvaluationResult = Evaluator->evaluate(); if (!LastEvaluationResult.hasValue()) { Ctx.emitError("Error evaluating model."); - return false; + return nullptr; } - int64_t Decision = *LastEvaluationResult->getTensorValue(0); - return static_cast(Decision); + return LastEvaluationResult->getTensorValue(0); } -int64_t ModelUnderTrainingRunner::getFeature(int Index) const { - return *Evaluator->getInput(Index); -} - -void ModelUnderTrainingRunner::setFeature(FeatureIndex Index, int64_t Value) { - size_t NumericIndex = static_cast(Index); - *(Evaluator->getInput(NumericIndex)) = Value; +void *ModelUnderTrainingRunner::getTensorUntyped(size_t Index) { + return Evaluator->getUntypedInput(Index); } std::unique_ptr llvm::getDevelopmentModeAdvisor( @@ -509,7 +493,8 @@ std::unique_ptr llvm::getDevelopmentModeAdvisor( ModelUnderTrainingRunner *MUTRPtr = nullptr; bool IsDoingInference = false; if (TFModelUnderTrainingPath.empty()) - Runner.reset(new NoInferenceModelRunner(Ctx)); + Runner.reset(new NoInferenceModelRunner( + Ctx, ModelUnderTrainingRunner::getInputFeatures())); else { auto MUTR = std::make_unique( Ctx, TFModelUnderTrainingPath); diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp index 6fc4c42bdd71..17f6fa503456 100644 --- a/llvm/lib/Analysis/MLInlineAdvisor.cpp +++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp @@ -245,29 +245,32 @@ std::unique_ptr MLInlineAdvisor::getAdviceImpl(CallBase &CB) { auto &CallerBefore = FAM.getResult(Caller); auto &CalleeBefore = FAM.getResult(Callee); - ModelRunner->setFeature(FeatureIndex::CalleeBasicBlockCount, - CalleeBefore.BasicBlockCount); - ModelRunner->setFeature(FeatureIndex::CallSiteHeight, - FunctionLevels[&Caller]); - ModelRunner->setFeature(FeatureIndex::NodeCount, NodeCount); - ModelRunner->setFeature(FeatureIndex::NrCtantParams, NrCtantParams); - ModelRunner->setFeature(FeatureIndex::EdgeCount, EdgeCount); - ModelRunner->setFeature(FeatureIndex::CallerUsers, CallerBefore.Uses); - ModelRunner->setFeature(FeatureIndex::CallerConditionallyExecutedBlocks, - CallerBefore.BlocksReachedFromConditionalInstruction); - ModelRunner->setFeature(FeatureIndex::CallerBasicBlockCount, - CallerBefore.BasicBlockCount); - ModelRunner->setFeature(FeatureIndex::CalleeConditionallyExecutedBlocks, - CalleeBefore.BlocksReachedFromConditionalInstruction); - ModelRunner->setFeature(FeatureIndex::CalleeUsers, CalleeBefore.Uses); - ModelRunner->setFeature(FeatureIndex::CostEstimate, CostEstimate); + *ModelRunner->getTensor(FeatureIndex::CalleeBasicBlockCount) = + CalleeBefore.BasicBlockCount; + *ModelRunner->getTensor(FeatureIndex::CallSiteHeight) = + FunctionLevels[&Caller]; + *ModelRunner->getTensor(FeatureIndex::NodeCount) = NodeCount; + *ModelRunner->getTensor(FeatureIndex::NrCtantParams) = NrCtantParams; + *ModelRunner->getTensor(FeatureIndex::EdgeCount) = EdgeCount; + *ModelRunner->getTensor(FeatureIndex::CallerUsers) = + CallerBefore.Uses; + *ModelRunner->getTensor( + FeatureIndex::CallerConditionallyExecutedBlocks) = + CallerBefore.BlocksReachedFromConditionalInstruction; + *ModelRunner->getTensor(FeatureIndex::CallerBasicBlockCount) = + CallerBefore.BasicBlockCount; + *ModelRunner->getTensor( + FeatureIndex::CalleeConditionallyExecutedBlocks) = + CalleeBefore.BlocksReachedFromConditionalInstruction; + *ModelRunner->getTensor(FeatureIndex::CalleeUsers) = + CalleeBefore.Uses; + *ModelRunner->getTensor(FeatureIndex::CostEstimate) = CostEstimate; // Add the cost features for (size_t I = 0; I < static_cast(InlineCostFeatureIndex::NumberOfFeatures); ++I) { - ModelRunner->setFeature( - inlineCostFeatureToMlFeature(static_cast(I)), - CostFeatures->at(I)); + *ModelRunner->getTensor(inlineCostFeatureToMlFeature( + static_cast(I))) = CostFeatures->at(I); } return getAdviceFromModel(CB, ORE); @@ -276,7 +279,8 @@ std::unique_ptr MLInlineAdvisor::getAdviceImpl(CallBase &CB) { std::unique_ptr MLInlineAdvisor::getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) { - return std::make_unique(this, CB, ORE, ModelRunner->run()); + return std::make_unique( + this, CB, ORE, static_cast(ModelRunner->evaluate())); } std::unique_ptr MLInlineAdvisor::getMandatoryAdvice(CallBase &CB, @@ -302,7 +306,8 @@ void MLInlineAdvice::reportContextForRemark( using namespace ore; OR << NV("Callee", Callee->getName()); for (size_t I = 0; I < NumberOfFeatures; ++I) - OR << NV(FeatureNameMap[I], getAdvisor()->getModelRunner().getFeature(I)); + OR << NV(FeatureNameMap[I], + *getAdvisor()->getModelRunner().getTensor(I)); OR << NV("ShouldInline", isInliningRecommended()); } diff --git a/llvm/lib/Analysis/NoInferenceModelRunner.cpp b/llvm/lib/Analysis/NoInferenceModelRunner.cpp new file mode 100644 index 000000000000..e8f90c22e818 --- /dev/null +++ b/llvm/lib/Analysis/NoInferenceModelRunner.cpp @@ -0,0 +1,33 @@ +//===- NoInferenceModelRunner.cpp - noop ML 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 +// +//===----------------------------------------------------------------------===// +// +// A pseudo model runner. We use it to store feature values when collecting +// logs for the default policy, in 'development' mode, but never ask it to +// 'run'. +//===----------------------------------------------------------------------===// +#include "llvm/Config/config.h" +#if defined(LLVM_HAVE_TF_API) + +#include "llvm/Analysis/NoInferenceModelRunner.h" +#include "llvm/Analysis/Utils/TFUtils.h" + +using namespace llvm; + +NoInferenceModelRunner::NoInferenceModelRunner( + LLVMContext &Ctx, const std::vector &Inputs) + : MLModelRunner(Ctx) { + ValuesBuffer.reserve(Inputs.size()); + for (const auto &TS : Inputs) + ValuesBuffer.push_back(std::make_unique(TS.getElementCount() * + TS.getElementByteSize())); +} + +void *NoInferenceModelRunner::getTensorUntyped(size_t Index) { + return ValuesBuffer[Index].get(); +} +#endif // defined(LLVM_HAVE_TF_API) \ No newline at end of file diff --git a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp index d2bf95388066..3d4c75dd4d84 100644 --- a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp +++ b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp @@ -35,12 +35,10 @@ 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: + void *evaluateUntyped() override; + void *getTensorUntyped(size_t Index) override; + std::vector FeatureIndices; int32_t ResultIndex = -1; std::unique_ptr CompiledModel; @@ -66,20 +64,14 @@ ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx) assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model"); } -int64_t ReleaseModeModelRunner::getFeature(int Index) const { - return *static_cast( +void *ReleaseModeModelRunner::getTensorUntyped(size_t Index) { + return reinterpret_cast( CompiledModel->arg_data(FeatureIndices[Index])); } -void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) { - *static_cast(CompiledModel->arg_data( - FeatureIndices[static_cast(Index)])) = Value; -} - -bool ReleaseModeModelRunner::run() { +void *ReleaseModeModelRunner::evaluateUntyped() { CompiledModel->Run(); - return static_cast( - *static_cast(CompiledModel->result_data(ResultIndex))); + return CompiledModel->result_data(ResultIndex); } std::unique_ptr diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt index 7e3e20e4af28..e40e30a15132 100644 --- a/llvm/unittests/Analysis/CMakeLists.txt +++ b/llvm/unittests/Analysis/CMakeLists.txt @@ -6,10 +6,11 @@ set(LLVM_LINK_COMPONENTS TransformUtils ) +set(MLGO_TESTS TFUtilsTest.cpp MLModelRunnerTest.cpp) if (DEFINED LLVM_HAVE_TF_API) - LIST(APPEND EXTRA_TESTS TFUtilsTest.cpp) + LIST(APPEND EXTRA_TESTS ${MLGO_TESTS}) else() - LIST(APPEND LLVM_OPTIONAL_SOURCES TFUtilsTest.cpp) + LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS}) endif() add_llvm_unittest_with_input_files(AnalysisTests diff --git a/llvm/unittests/Analysis/MLModelRunnerTest.cpp b/llvm/unittests/Analysis/MLModelRunnerTest.cpp new file mode 100644 index 000000000000..9794365ca51c --- /dev/null +++ b/llvm/unittests/Analysis/MLModelRunnerTest.cpp @@ -0,0 +1,33 @@ +//===- MLModelRunnerTest.cpp - test for MLModelRunner ---------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/NoInferenceModelRunner.h" +#include "gtest/gtest.h" + +using namespace llvm; + +TEST(NoInferenceModelRunner, AccessTensors) { + const std::vector Inputs{ + TensorSpec::createSpec("F1", {1}), + TensorSpec::createSpec("F2", {10}), + TensorSpec::createSpec("F2", {5}), + }; + LLVMContext Ctx; + NoInferenceModelRunner NIMR(Ctx, Inputs); + NIMR.getTensor(0)[0] = 1; + std::memcpy(NIMR.getTensor(1), + std::vector{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(), + 10 * sizeof(int64_t)); + std::memcpy(NIMR.getTensor(2), + std::vector{0.1, 0.2, 0.3, 0.4, 0.5}.data(), + 5 * sizeof(float)); + ASSERT_EQ(NIMR.getTensor(0)[0], 1); + ASSERT_EQ(NIMR.getTensor(1)[8], 9); + ASSERT_EQ(NIMR.getTensor(2)[1], 0.2f); +} \ No newline at end of file