diff --git a/llvm/include/llvm/Analysis/MLModelRunner.h b/llvm/include/llvm/Analysis/MLModelRunner.h index 7cfa6efedf108..2281b63448354 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 @@ class MLModelRunner { 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 0000000000000..6f6f5a009b1c7 --- /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 1f6be0e60eb99..e4a4b7428cb6c 100644 --- a/llvm/include/llvm/Analysis/Utils/TFUtils.h +++ b/llvm/include/llvm/Analysis/Utils/TFUtils.h @@ -246,8 +246,10 @@ class TFModelEvaluator final { /// 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 9da07cb1c4850..9dfc619f2b135 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 d87fa849d839a..b76d63dddeb01 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 @@ class LoggingMLInlineAdvice : public MLInlineAdvice { 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 @@ class ModelUnderTrainingRunner final : public MLModelRunner { 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); -} - -int64_t ModelUnderTrainingRunner::getFeature(int Index) const { - return *Evaluator->getInput(Index); + return LastEvaluationResult->getTensorValue(0); } -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 6fc4c42bdd718..17f6fa5034567 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 0000000000000..e8f90c22e818e --- /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 d2bf95388066b..3d4c75dd4d843 100644 --- a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp +++ b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp @@ -35,12 +35,10 @@ class ReleaseModeModelRunner final : public MLModelRunner { 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 7e3e20e4af287..e40e30a15132f 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 0000000000000..9794365ca51c6 --- /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