[llvm] [mlgo] Support composite AOT-ed models (PR #96276)
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Mon Jun 24 11:09:23 PDT 2024
https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/96276
>From 35e6530844ab790f6eadcc1cdd38eca4cd8c7c70 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Thu, 20 Jun 2024 22:06:11 -0700
Subject: [PATCH 1/4] [mlgo] Support composite AOT-ed models
This adds support for multiple models for the same agent, and allow the
user select one via a command line flag, in the AOT case, where we embed
the models in the compiler.
To avoid build setup complexity, the support is delegated to the saved
model. Since saved models define computational graphs, we can generate a
composite model (prior to building and embedding it in LLVM) that exposes
an extra feature with a predefined name - `_model_selector`. The model,
then, delegates internally to contained models based on that feature value.
If the model doesn't expose such a feature but the user passes one, we
report error.
If the model exposes such a feature but the user doesn't pass one, we
expect the model to pick a default.
Invalid model selector values are expected to be handled by the saved
model.
The model is selected using a string name by the LLVM user, but the model
uses a pair of uint64 values - the high and low of the MD5 hash of the
name.
A tool composing models would, then, just need to:
- expose the extra feature, `_model_selector`, shape (2,), uint64 data type
- test its value against the MD5 hash, in the [high, low] order, of
contained models based on a user-specified name (which the user will
then use as flag value to the compiler)
Agents like the inline advisor or the regalloc eviction advisor just need
to add a flag to capture the name of an agent and pass it to
`ReleaseModeModelRunner` at construction.
---
.../llvm/Analysis/ReleaseModeModelRunner.h | 91 ++++++++++--
llvm/lib/Analysis/MLInlineAdvisor.cpp | 6 +-
llvm/unittests/Analysis/MLModelRunnerTest.cpp | 134 ++++++++++++++++--
3 files changed, 204 insertions(+), 27 deletions(-)
diff --git a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
index 91855138fe18e..629d62295dafb 100644
--- a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
+++ b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
@@ -17,14 +17,39 @@
#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/Analysis/TensorSpec.h"
#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MD5.h"
#include <memory>
-#include <vector>
namespace llvm {
/// ReleaseModeModelRunner - production mode implementation of the
/// MLModelRunner. It uses an AOT-compiled SavedModel for efficient execution.
+struct EmbeddedModelRunnerOptions {
+ /// Feed and Fetch feature prefixes - i.e. a feature named "foo" will be
+ /// looked up as {FeedPrefix}_foo; and the output named "bar" will be looked
+ /// up as {FetchPrefix}_bar
+ StringRef FeedPrefix = "feed_";
+ StringRef FetchPrefix = "fetch_";
+
+ /// ModelSelector is the name (recognized by the AOT-ed model) of a sub-model
+ /// to use. "" is an alias for a default model to select.
+ StringRef ModelSelector = "";
+
+ EmbeddedModelRunnerOptions &setFeedPrefix(StringRef Value) {
+ FeedPrefix = Value;
+ return *this;
+ }
+ EmbeddedModelRunnerOptions &setFetchPrefix(StringRef Value) {
+ FetchPrefix = Value;
+ return *this;
+ }
+ EmbeddedModelRunnerOptions &setModelSelector(StringRef Value) {
+ ModelSelector = Value;
+ return *this;
+ }
+};
+
template <class TGen>
class ReleaseModeModelRunner final : public MLModelRunner {
public:
@@ -32,22 +57,49 @@ class ReleaseModeModelRunner final : public MLModelRunner {
/// std::array or std::vector, that has a size() method.
template <class FType>
ReleaseModeModelRunner(LLVMContext &Ctx, const FType &InputSpec,
- StringRef DecisionName, StringRef FeedPrefix = "feed_",
- StringRef FetchPrefix = "fetch_")
- : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()),
+ StringRef DecisionName,
+ const EmbeddedModelRunnerOptions &Options = {})
+ : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size() + 1),
CompiledModel(std::make_unique<TGen>()) {
assert(CompiledModel && "The CompiledModel should be valid");
-
- for (size_t I = 0; I < InputSpec.size(); ++I) {
- const int Index =
- CompiledModel->LookupArgIndex(FeedPrefix.str() + InputSpec[I].name());
- void *Buffer = nullptr;
- if (Index >= 0)
- Buffer = CompiledModel->arg_data(Index);
- setUpBufferForTensor(I, InputSpec[I], Buffer);
+ // Set up the model_selector past all the InputSpecs in all cases.
+ // - if the model doesn't have such a feature, but the user requested it,
+ // we report error
+ // - otherwise, if the user doesn't specify it, we set the selector to {0,
+ // 0} which the model should handle as a "default".
+ // - finally, we compute the MD5 hash of the user input and set the value
+ // of the model selector to {high, low}
+ bool InputIsPresent = true;
+ populateTensor(InputSpec.size(),
+ TensorSpec::createSpec<uint64_t>("_model_selector", {2}),
+ Options.FeedPrefix, InputIsPresent);
+ uint64_t High = 0;
+ uint64_t Low = 0;
+ if (!Options.ModelSelector.empty()) {
+ if (!InputIsPresent) {
+ Ctx.emitError("A model selector was specified but the underlying model "
+ "does not expose a _model_selector input");
+ // continue with the set up in case there's some custom diagnostics
+ // handler installed. Since the model has no idea what a model selector
+ // input feature is, it'll be ignored. Otherwise we exited above.
+ }
+ const auto Hash = MD5::hash(
+ {reinterpret_cast<const uint8_t *>(Options.ModelSelector.data()),
+ Options.ModelSelector.size()});
+
+ High = Hash.high();
+ Low = Hash.low();
}
-
- ResultIndex = CompiledModel->LookupResultIndex(FetchPrefix.str() +
+ getTensor<uint64_t>(InputSpec.size())[0] = High;
+ getTensor<uint64_t>(InputSpec.size())[1] = Low;
+ // At this point, the model selector is set up. If the user didn't provide
+ // one, it'll be (0, 0) which the composite model should treat as "default"
+ // - whichever that is. If the model isn't composite, the rest of the set up
+ // continues normally, but an error would be emitted.
+ for (size_t I = 0; I < InputSpec.size(); ++I)
+ populateTensor(I, InputSpec[I], Options.FeedPrefix, InputIsPresent);
+
+ ResultIndex = CompiledModel->LookupResultIndex(Options.FetchPrefix.str() +
DecisionName.str());
assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model");
}
@@ -59,6 +111,17 @@ class ReleaseModeModelRunner final : public MLModelRunner {
}
private:
+ void populateTensor(size_t Pos, const TensorSpec &Spec, StringRef Prefix,
+ bool &InputIsPresent) {
+ const int Index =
+ CompiledModel->LookupArgIndex((Prefix + Spec.name()).str());
+ void *Buffer = nullptr;
+ InputIsPresent = Index >= 0;
+ if (InputIsPresent)
+ Buffer = CompiledModel->arg_data(Index);
+ setUpBufferForTensor(Pos, Spec, Buffer);
+ }
+
void *evaluateUntyped() override {
CompiledModel->Run();
return CompiledModel->result_data(ResultIndex);
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 21946572339b9..2c389d8825892 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -56,6 +56,9 @@ static cl::opt<SkipMLPolicyCriteria> SkipPolicy(
clEnumValN(SkipMLPolicyCriteria::IfCallerIsNotCold,
"if-caller-not-cold", "if the caller is not cold")));
+static cl::opt<std::string> ModelSelector("ml-inliner-model-selector",
+ cl::Hidden, cl::init(""));
+
#if defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
// codegen-ed file
#include "InlinerSizeModel.h" // NOLINT
@@ -73,7 +76,8 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
std::unique_ptr<MLModelRunner> AOTRunner;
if (InteractiveChannelBaseName.empty())
AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
- M.getContext(), FeatureMap, DecisionName);
+ M.getContext(), FeatureMap, DecisionName,
+ EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
else {
auto Features = FeatureMap;
if (InteractiveIncludeDefault)
diff --git a/llvm/unittests/Analysis/MLModelRunnerTest.cpp b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
index 007a8cfef0433..9ad046059fdc3 100644
--- a/llvm/unittests/Analysis/MLModelRunnerTest.cpp
+++ b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
@@ -7,10 +7,12 @@
//===----------------------------------------------------------------------===//
#include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/ADT/StringExtras.h"
#include "llvm/Analysis/InteractiveModelRunner.h"
#include "llvm/Analysis/NoInferenceModelRunner.h"
#include "llvm/Analysis/ReleaseModeModelRunner.h"
#include "llvm/Support/BinaryByteStream.h"
+#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/JSON.h"
@@ -28,14 +30,17 @@ namespace llvm {
// This is a mock of the kind of AOT-generated model evaluator. It has 2 tensors
// of shape {1}, and 'evaluation' adds them.
// The interface is the one expected by ReleaseModelRunner.
-class MockAOTModel final {
+class MockAOTModelBase {
+protected:
int64_t A = 0;
int64_t B = 0;
int64_t R = 0;
public:
- MockAOTModel() = default;
- int LookupArgIndex(const std::string &Name) {
+ MockAOTModelBase() = default;
+ virtual ~MockAOTModelBase() = default;
+
+ virtual int LookupArgIndex(const std::string &Name) {
if (Name == "prefix_a")
return 0;
if (Name == "prefix_b")
@@ -43,13 +48,13 @@ class MockAOTModel final {
return -1;
}
int LookupResultIndex(const std::string &) { return 0; }
- void Run() { R = A + B; }
- void *result_data(int RIndex) {
+ virtual void Run() = 0;
+ virtual void *result_data(int RIndex) {
if (RIndex == 0)
return &R;
return nullptr;
}
- void *arg_data(int Index) {
+ virtual void *arg_data(int Index) {
switch (Index) {
case 0:
return &A;
@@ -60,6 +65,64 @@ class MockAOTModel final {
}
}
};
+
+class AdditionAOTModel final : public MockAOTModelBase {
+public:
+ AdditionAOTModel() = default;
+ void Run() override { R = A + B; }
+};
+
+class DiffAOTModel final : public MockAOTModelBase {
+public:
+ DiffAOTModel() = default;
+ void Run() override { R = A - B; }
+};
+
+static const char *M1Selector = "the model that subtracts";
+static const char *M2Selector = "the model that adds";
+
+static auto Hash1 = MD5::hash(arrayRefFromStringRef(M1Selector));
+static auto Hash2 = MD5::hash(arrayRefFromStringRef(M2Selector));
+class ComposedAOTModel final {
+ DiffAOTModel M1;
+ AdditionAOTModel M2;
+ uint64_t Selector[2] = {0};
+
+ bool isHashSameAsSelector(const std::pair<uint64_t, uint64_t> &Words) const {
+ return Selector[0] == Words.first && Selector[1] == Words.second;
+ }
+ MockAOTModelBase *getModel() {
+ if (isHashSameAsSelector(Hash1.words()) || (!Selector[0] && !Selector[1]))
+ return &M1;
+ if (isHashSameAsSelector(Hash2.words()))
+ return &M2;
+ llvm_unreachable("Should be one of the two");
+ }
+
+public:
+ ComposedAOTModel() = default;
+ int LookupArgIndex(const std::string &Name) {
+ if (Name == "prefix__model_selector")
+ return 2;
+ return getModel()->LookupArgIndex(Name);
+ }
+ int LookupResultIndex(const std::string &Name) {
+ return getModel()->LookupResultIndex(Name);
+ }
+ void *arg_data(int Index) {
+ if (Index == 2)
+ return Selector;
+ return getModel()->arg_data(Index);
+ }
+ void *result_data(int RIndex) { return getModel()->result_data(RIndex); }
+ void Run() { getModel()->Run(); }
+};
+
+static EmbeddedModelRunnerOptions makeOptions() {
+ EmbeddedModelRunnerOptions Opts;
+ Opts.setFeedPrefix("prefix_");
+ return Opts;
+}
} // namespace llvm
TEST(NoInferenceModelRunner, AccessTensors) {
@@ -86,8 +149,8 @@ TEST(ReleaseModeRunner, NormalUse) {
LLVMContext Ctx;
std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
TensorSpec::createSpec<int64_t>("b", {1})};
- auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
- Ctx, Inputs, "", "prefix_");
+ auto Evaluator = std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(
+ Ctx, Inputs, "", makeOptions());
*Evaluator->getTensor<int64_t>(0) = 1;
*Evaluator->getTensor<int64_t>(1) = 2;
EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
@@ -100,8 +163,8 @@ TEST(ReleaseModeRunner, ExtraFeatures) {
std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
TensorSpec::createSpec<int64_t>("b", {1}),
TensorSpec::createSpec<int64_t>("c", {1})};
- auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
- Ctx, Inputs, "", "prefix_");
+ auto Evaluator = std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(
+ Ctx, Inputs, "", makeOptions());
*Evaluator->getTensor<int64_t>(0) = 1;
*Evaluator->getTensor<int64_t>(1) = 2;
*Evaluator->getTensor<int64_t>(2) = -3;
@@ -118,8 +181,8 @@ TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {
TensorSpec::createSpec<int64_t>("c", {1}),
TensorSpec::createSpec<int64_t>("b", {1}),
};
- auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
- Ctx, Inputs, "", "prefix_");
+ auto Evaluator = std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(
+ Ctx, Inputs, "", makeOptions());
*Evaluator->getTensor<int64_t>(0) = 1; // a
*Evaluator->getTensor<int64_t>(1) = 2; // c
*Evaluator->getTensor<int64_t>(2) = -3; // b
@@ -129,6 +192,53 @@ TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {
EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
}
+// We expect an error to be reported early if the user tried to specify a model
+// selector, but the model in fact doesn't support that.
+TEST(ReleaseModelRunner, ModelSelectorNoInput) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1})};
+ EXPECT_DEATH(std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(
+ Ctx, Inputs, "", makeOptions().setModelSelector(M2Selector)),
+ "A model selector was specified but the underlying model does "
+ "not expose a _model_selector input");
+}
+
+// Test that we correctly set up the _model_selector tensor value. We are only
+// responsbile for what happens if the user doesn't specify a value (but the
+// model supports the feature), or if the user specifies one, and we correctly
+// populate the tensor, and do so upfront (in case the model implementation
+// needs that for subsequent tensor buffer lookups).
+TEST(ReleaseModelRunner, ModelSelector) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1})};
+ auto Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
+ Ctx, Inputs, "", makeOptions());
+ *Evaluator->getTensor<int64_t>(0) = 1;
+ *Evaluator->getTensor<int64_t>(1) = 2;
+ // No selector was specified, so we'll use the default one, which in this case
+ // is M1 in ComposedAOTModel's implementation.
+ EXPECT_EQ(Evaluator->evaluate<int64_t>(), -1);
+
+ // This explicitly asks for M1
+ Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
+ Ctx, Inputs, "", makeOptions().setModelSelector(M1Selector));
+ *Evaluator->getTensor<int64_t>(0) = 1;
+ *Evaluator->getTensor<int64_t>(1) = 2;
+ EXPECT_EQ(Evaluator->evaluate<int64_t>(), -1);
+
+ // Ask for M2
+ Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
+ Ctx, Inputs, "", makeOptions().setModelSelector(M2Selector));
+ *Evaluator->getTensor<int64_t>(0) = 1;
+ *Evaluator->getTensor<int64_t>(1) = 2;
+ EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
+
+ // Asking for a model that's not supported isn't handled by our infra and we
+ // expect the model implementation to fail at a point.
+}
+
#if defined(LLVM_ON_UNIX)
TEST(InteractiveModelRunner, Evaluation) {
LLVMContext Ctx;
>From 8551dbf91cb9689bfdc7e65dd97c138cecfa8738 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Fri, 21 Jun 2024 08:51:11 -0700
Subject: [PATCH 2/4] changed how "" selector is handled, to error
---
.../llvm/Analysis/ReleaseModeModelRunner.h | 29 +++++++++++--------
llvm/unittests/Analysis/MLModelRunnerTest.cpp | 25 +++++++++-------
2 files changed, 31 insertions(+), 23 deletions(-)
diff --git a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
index 629d62295dafb..f8c963e4f41cb 100644
--- a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
+++ b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
@@ -33,7 +33,7 @@ struct EmbeddedModelRunnerOptions {
StringRef FetchPrefix = "fetch_";
/// ModelSelector is the name (recognized by the AOT-ed model) of a sub-model
- /// to use. "" is an alias for a default model to select.
+ /// to use. "" is allowed if the model doesn't support sub-models.
StringRef ModelSelector = "";
EmbeddedModelRunnerOptions &setFeedPrefix(StringRef Value) {
@@ -64,25 +64,28 @@ class ReleaseModeModelRunner final : public MLModelRunner {
assert(CompiledModel && "The CompiledModel should be valid");
// Set up the model_selector past all the InputSpecs in all cases.
// - if the model doesn't have such a feature, but the user requested it,
- // we report error
- // - otherwise, if the user doesn't specify it, we set the selector to {0,
- // 0} which the model should handle as a "default".
+ // we report error. Same if the model supports it but the user didn't
+ // specify it
// - finally, we compute the MD5 hash of the user input and set the value
// of the model selector to {high, low}
bool InputIsPresent = true;
populateTensor(InputSpec.size(),
TensorSpec::createSpec<uint64_t>("_model_selector", {2}),
Options.FeedPrefix, InputIsPresent);
+
+ // If we hit the "report an error" cases outlined above, continue with the
+ // set up in case there's some custom diagnostics handler installed and it
+ // doesn't promptly exit.
+ if (Options.ModelSelector.empty() && InputIsPresent)
+ Ctx.emitError(
+ "A model selector was not specified but the underlying model "
+ "requires selecting one because it exposes a _model_selector input");
uint64_t High = 0;
uint64_t Low = 0;
if (!Options.ModelSelector.empty()) {
- if (!InputIsPresent) {
+ if (!InputIsPresent)
Ctx.emitError("A model selector was specified but the underlying model "
"does not expose a _model_selector input");
- // continue with the set up in case there's some custom diagnostics
- // handler installed. Since the model has no idea what a model selector
- // input feature is, it'll be ignored. Otherwise we exited above.
- }
const auto Hash = MD5::hash(
{reinterpret_cast<const uint8_t *>(Options.ModelSelector.data()),
Options.ModelSelector.size()});
@@ -93,9 +96,8 @@ class ReleaseModeModelRunner final : public MLModelRunner {
getTensor<uint64_t>(InputSpec.size())[0] = High;
getTensor<uint64_t>(InputSpec.size())[1] = Low;
// At this point, the model selector is set up. If the user didn't provide
- // one, it'll be (0, 0) which the composite model should treat as "default"
- // - whichever that is. If the model isn't composite, the rest of the set up
- // continues normally, but an error would be emitted.
+ // one, but the model has a _model_selector, it'll be set to (0, 0) which
+ // the composite model should treat as error as part of its implementation.
for (size_t I = 0; I < InputSpec.size(); ++I)
populateTensor(I, InputSpec[I], Options.FeedPrefix, InputIsPresent);
@@ -111,6 +113,9 @@ class ReleaseModeModelRunner final : public MLModelRunner {
}
private:
+ // fetch the model-provided buffer for the given Spec, or let MLModelRunner
+ // create a scratch buffer. Indicate back to the caller if the model had that
+ // input in the first place.
void populateTensor(size_t Pos, const TensorSpec &Spec, StringRef Prefix,
bool &InputIsPresent) {
const int Index =
diff --git a/llvm/unittests/Analysis/MLModelRunnerTest.cpp b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
index 9ad046059fdc3..0c18c251c0bf8 100644
--- a/llvm/unittests/Analysis/MLModelRunnerTest.cpp
+++ b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
@@ -92,7 +92,7 @@ class ComposedAOTModel final {
return Selector[0] == Words.first && Selector[1] == Words.second;
}
MockAOTModelBase *getModel() {
- if (isHashSameAsSelector(Hash1.words()) || (!Selector[0] && !Selector[1]))
+ if (isHashSameAsSelector(Hash1.words()))
return &M1;
if (isHashSameAsSelector(Hash2.words()))
return &M2;
@@ -194,7 +194,7 @@ TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {
// We expect an error to be reported early if the user tried to specify a model
// selector, but the model in fact doesn't support that.
-TEST(ReleaseModelRunner, ModelSelectorNoInput) {
+TEST(ReleaseModelRunner, ModelSelectorNoInputFeaturePresent) {
LLVMContext Ctx;
std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
TensorSpec::createSpec<int64_t>("b", {1})};
@@ -204,6 +204,17 @@ TEST(ReleaseModelRunner, ModelSelectorNoInput) {
"not expose a _model_selector input");
}
+TEST(ReleaseModelRunner, ModelSelectorNoSelectorGiven) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1})};
+ EXPECT_DEATH(
+ std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
+ Ctx, Inputs, "", makeOptions()),
+ "A model selector was not specified but the underlying model requires "
+ "selecting one because it exposes a _model_selector input");
+}
+
// Test that we correctly set up the _model_selector tensor value. We are only
// responsbile for what happens if the user doesn't specify a value (but the
// model supports the feature), or if the user specifies one, and we correctly
@@ -213,16 +224,8 @@ TEST(ReleaseModelRunner, ModelSelector) {
LLVMContext Ctx;
std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
TensorSpec::createSpec<int64_t>("b", {1})};
- auto Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
- Ctx, Inputs, "", makeOptions());
- *Evaluator->getTensor<int64_t>(0) = 1;
- *Evaluator->getTensor<int64_t>(1) = 2;
- // No selector was specified, so we'll use the default one, which in this case
- // is M1 in ComposedAOTModel's implementation.
- EXPECT_EQ(Evaluator->evaluate<int64_t>(), -1);
-
// This explicitly asks for M1
- Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
+ auto Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(
Ctx, Inputs, "", makeOptions().setModelSelector(M1Selector));
*Evaluator->getTensor<int64_t>(0) = 1;
*Evaluator->getTensor<int64_t>(1) = 2;
>From 377a825b8de30b0c5a38deba99441aca2e2f1f70 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 24 Jun 2024 09:37:03 -0700
Subject: [PATCH 3/4] feedback
---
llvm/include/llvm/Analysis/ReleaseModeModelRunner.h | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
index f8c963e4f41cb..9fb4ff4ec4013 100644
--- a/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
+++ b/llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
@@ -14,6 +14,7 @@
#ifndef LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
#define LLVM_ANALYSIS_RELEASEMODEMODELRUNNER_H
+#include "llvm/ADT/StringExtras.h"
#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/Analysis/TensorSpec.h"
#include "llvm/Support/ErrorHandling.h"
@@ -86,10 +87,7 @@ class ReleaseModeModelRunner final : public MLModelRunner {
if (!InputIsPresent)
Ctx.emitError("A model selector was specified but the underlying model "
"does not expose a _model_selector input");
- const auto Hash = MD5::hash(
- {reinterpret_cast<const uint8_t *>(Options.ModelSelector.data()),
- Options.ModelSelector.size()});
-
+ const auto Hash = MD5::hash(arrayRefFromStringRef(Options.ModelSelector));
High = Hash.high();
Low = Hash.low();
}
@@ -97,7 +95,9 @@ class ReleaseModeModelRunner final : public MLModelRunner {
getTensor<uint64_t>(InputSpec.size())[1] = Low;
// At this point, the model selector is set up. If the user didn't provide
// one, but the model has a _model_selector, it'll be set to (0, 0) which
- // the composite model should treat as error as part of its implementation.
+ // the composite model should treat as error as part of its implementation
+ // (but that should only matter if there is a custom handler that doesn't
+ // exit on error)
for (size_t I = 0; I < InputSpec.size(); ++I)
populateTensor(I, InputSpec[I], Options.FeedPrefix, InputIsPresent);
>From d1648b623521fd78ae2599168c9308f88e8a126a Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 24 Jun 2024 11:08:51 -0700
Subject: [PATCH 4/4] auto -> MD5::MD5Result
---
llvm/unittests/Analysis/MLModelRunnerTest.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/llvm/unittests/Analysis/MLModelRunnerTest.cpp b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
index 0c18c251c0bf8..53ad2fc5c4b08 100644
--- a/llvm/unittests/Analysis/MLModelRunnerTest.cpp
+++ b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
@@ -81,8 +81,8 @@ class DiffAOTModel final : public MockAOTModelBase {
static const char *M1Selector = "the model that subtracts";
static const char *M2Selector = "the model that adds";
-static auto Hash1 = MD5::hash(arrayRefFromStringRef(M1Selector));
-static auto Hash2 = MD5::hash(arrayRefFromStringRef(M2Selector));
+static MD5::MD5Result Hash1 = MD5::hash(arrayRefFromStringRef(M1Selector));
+static MD5::MD5Result Hash2 = MD5::hash(arrayRefFromStringRef(M2Selector));
class ComposedAOTModel final {
DiffAOTModel M1;
AdditionAOTModel M2;
More information about the llvm-commits
mailing list