[llvm] 059e034 - [NFC][mlgo] Generalize model runner interface
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Wed Dec 8 20:11:10 PST 2021
Author: Mircea Trofin
Date: 2021-12-08T20:10:58-08:00
New Revision: 059e03476cbbb7349b75fa9cc26bc9da0d491529
URL: https://github.com/llvm/llvm-project/commit/059e03476cbbb7349b75fa9cc26bc9da0d491529
DIFF: https://github.com/llvm/llvm-project/commit/059e03476cbbb7349b75fa9cc26bc9da0d491529.diff
LOG: [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
Added:
llvm/include/llvm/Analysis/NoInferenceModelRunner.h
llvm/lib/Analysis/NoInferenceModelRunner.cpp
llvm/unittests/Analysis/MLModelRunnerTest.cpp
Modified:
llvm/include/llvm/Analysis/MLModelRunner.h
llvm/include/llvm/Analysis/Utils/TFUtils.h
llvm/lib/Analysis/CMakeLists.txt
llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
llvm/lib/Analysis/MLInlineAdvisor.cpp
llvm/lib/Analysis/ReleaseModeModelRunner.cpp
llvm/unittests/Analysis/CMakeLists.txt
Removed:
################################################################################
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 @@ 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 <typename T> T evaluate() {
+ return *reinterpret_cast<T *>(evaluateUntyped());
+ }
+
+ template <typename T, typename I> T *getTensor(I FeatureID) {
+ return reinterpret_cast<T *>(
+ getTensorUntyped(static_cast<size_t>(FeatureID)));
+ }
+
+ template <typename T, typename I> const T *getTensor(I FeatureID) const {
+ return reinterpret_cast<const T *>(
+ getTensorUntyped(static_cast<size_t>(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<MLModelRunner *>(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<TensorSpec> &Inputs);
+
+private:
+ void *evaluateUntyped() override {
+ llvm_unreachable("We shouldn't call run on this model runner.");
+ }
+ void *getTensorUntyped(size_t Index) override;
+
+ std::vector<std::unique_ptr<char[]>> 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 @@ class TFModelEvaluator final {
/// otherwise.
bool isValid() const { return !!Impl; }
-private:
+ /// Untyped access to input.
void *getUntypedInput(size_t Index);
+
+private:
std::unique_ptr<TFModelEvaluatorImpl> 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 @@ 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<int>(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<LoggedFeatureSpec> &outputLoggedFeatureSpecs() const {
@@ -308,18 +286,31 @@ class ModelUnderTrainingRunner final : public MLModelRunner {
return LastEvaluationResult;
}
+ static const std::vector<TensorSpec> getInputFeatures() {
+ std::vector<TensorSpec> InputSpecs;
+ for (size_t I = 0; I < NumberOfFeatures; ++I)
+ InputSpecs.push_back(TensorSpec::createSpec<int64_t>(
+ TFFeedPrefix + FeatureNameMap[I], {1}));
+ append_range(InputSpecs, TrainingOnlyFeatures);
+ return InputSpecs;
+ }
+
private:
std::unique_ptr<TFModelEvaluator> Evaluator;
std::vector<LoggedFeatureSpec> OutputSpecs;
Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
+ void *evaluateUntyped() override;
+ void *getTensorUntyped(size_t Index) override;
// The training framework needs some additional features.
- const std::vector<TensorSpec> TrainingOnlyFeatures{
- TensorSpec::createSpec<int64_t>(TFFeedPrefix + "inlining_default", {1}),
- TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),
- TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),
- TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})};
+ const static std::vector<TensorSpec> TrainingOnlyFeatures;
};
+
+const std::vector<TensorSpec> ModelUnderTrainingRunner::TrainingOnlyFeatures{
+ TensorSpec::createSpec<int64_t>(TFFeedPrefix + "inlining_default", {1}),
+ TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),
+ TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),
+ TensorSpec::createSpec<int32_t>(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<int64_t>(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<bool>(ModelRunner->evaluate<int64_t>())
+ : DefaultAdvice;
return std::make_unique<LoggingMLInlineAdvice>(
/*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<TensorSpec> InputSpecs;
- for (size_t I = 0; I < NumberOfFeatures; ++I)
- InputSpecs.push_back(
- TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
- append_range(InputSpecs, TrainingOnlyFeatures);
+ std::vector<TensorSpec> 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<int64_t>(0);
- return static_cast<bool>(Decision);
-}
-
-int64_t ModelUnderTrainingRunner::getFeature(int Index) const {
- return *Evaluator->getInput<int64_t>(Index);
+ return LastEvaluationResult->getTensorValue<int64_t>(0);
}
-void ModelUnderTrainingRunner::setFeature(FeatureIndex Index, int64_t Value) {
- size_t NumericIndex = static_cast<size_t>(Index);
- *(Evaluator->getInput<int64_t>(NumericIndex)) = Value;
+void *ModelUnderTrainingRunner::getTensorUntyped(size_t Index) {
+ return Evaluator->getUntypedInput(Index);
}
std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
@@ -509,7 +493,8 @@ std::unique_ptr<InlineAdvisor> 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<ModelUnderTrainingRunner>(
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<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
auto &CallerBefore = FAM.getResult<FunctionPropertiesAnalysis>(Caller);
auto &CalleeBefore = FAM.getResult<FunctionPropertiesAnalysis>(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<int64_t>(FeatureIndex::CalleeBasicBlockCount) =
+ CalleeBefore.BasicBlockCount;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::CallSiteHeight) =
+ FunctionLevels[&Caller];
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::NodeCount) = NodeCount;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::NrCtantParams) = NrCtantParams;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::EdgeCount) = EdgeCount;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::CallerUsers) =
+ CallerBefore.Uses;
+ *ModelRunner->getTensor<int64_t>(
+ FeatureIndex::CallerConditionallyExecutedBlocks) =
+ CallerBefore.BlocksReachedFromConditionalInstruction;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::CallerBasicBlockCount) =
+ CallerBefore.BasicBlockCount;
+ *ModelRunner->getTensor<int64_t>(
+ FeatureIndex::CalleeConditionallyExecutedBlocks) =
+ CalleeBefore.BlocksReachedFromConditionalInstruction;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::CalleeUsers) =
+ CalleeBefore.Uses;
+ *ModelRunner->getTensor<int64_t>(FeatureIndex::CostEstimate) = CostEstimate;
// Add the cost features
for (size_t I = 0;
I < static_cast<size_t>(InlineCostFeatureIndex::NumberOfFeatures); ++I) {
- ModelRunner->setFeature(
- inlineCostFeatureToMlFeature(static_cast<InlineCostFeatureIndex>(I)),
- CostFeatures->at(I));
+ *ModelRunner->getTensor<int64_t>(inlineCostFeatureToMlFeature(
+ static_cast<InlineCostFeatureIndex>(I))) = CostFeatures->at(I);
}
return getAdviceFromModel(CB, ORE);
@@ -276,7 +279,8 @@ std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
std::unique_ptr<MLInlineAdvice>
MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
OptimizationRemarkEmitter &ORE) {
- return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run());
+ return std::make_unique<MLInlineAdvice>(
+ this, CB, ORE, static_cast<bool>(ModelRunner->evaluate<int64_t>()));
}
std::unique_ptr<InlineAdvice> 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<int64_t>(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<TensorSpec> &Inputs)
+ : MLModelRunner(Ctx) {
+ ValuesBuffer.reserve(Inputs.size());
+ for (const auto &TS : Inputs)
+ ValuesBuffer.push_back(std::make_unique<char[]>(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 @@ 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<int32_t> FeatureIndices;
int32_t ResultIndex = -1;
std::unique_ptr<llvm::InlinerSizeModel> 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<int64_t *>(
+void *ReleaseModeModelRunner::getTensorUntyped(size_t Index) {
+ return reinterpret_cast<char *>(
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() {
+void *ReleaseModeModelRunner::evaluateUntyped() {
CompiledModel->Run();
- return static_cast<bool>(
- *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
+ return CompiledModel->result_data(ResultIndex);
}
std::unique_ptr<InlineAdvisor>
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<TensorSpec> Inputs{
+ TensorSpec::createSpec<int64_t>("F1", {1}),
+ TensorSpec::createSpec<int64_t>("F2", {10}),
+ TensorSpec::createSpec<float>("F2", {5}),
+ };
+ LLVMContext Ctx;
+ NoInferenceModelRunner NIMR(Ctx, Inputs);
+ NIMR.getTensor<int64_t>(0)[0] = 1;
+ std::memcpy(NIMR.getTensor<int64_t>(1),
+ std::vector<int64_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(),
+ 10 * sizeof(int64_t));
+ std::memcpy(NIMR.getTensor<float>(2),
+ std::vector<float>{0.1, 0.2, 0.3, 0.4, 0.5}.data(),
+ 5 * sizeof(float));
+ ASSERT_EQ(NIMR.getTensor<int64_t>(0)[0], 1);
+ ASSERT_EQ(NIMR.getTensor<int64_t>(1)[8], 9);
+ ASSERT_EQ(NIMR.getTensor<float>(2)[1], 0.2f);
+}
\ No newline at end of file
More information about the llvm-commits
mailing list