390 lines · cpp
1//===- MLModelRunnerTest.cpp - test for MLModelRunner ---------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/Analysis/MLModelRunner.h"10#include "llvm/ADT/StringExtras.h"11#include "llvm/Analysis/InteractiveModelRunner.h"12#include "llvm/Analysis/NoInferenceModelRunner.h"13#include "llvm/Analysis/ReleaseModeModelRunner.h"14#include "llvm/Config/llvm-config.h" // for LLVM_ON_UNIX15#include "llvm/Support/BinaryByteStream.h"16#include "llvm/Support/ErrorHandling.h"17#include "llvm/Support/FileSystem.h"18#include "llvm/Support/FileUtilities.h"19#include "llvm/Support/JSON.h"20#include "llvm/Support/Path.h"21#include "llvm/Support/raw_ostream.h"22#include "llvm/Testing/Support/SupportHelpers.h"23#include "gtest/gtest.h"24#include <atomic>25#include <thread>26 27using namespace llvm;28 29namespace llvm {30// This is a mock of the kind of AOT-generated model evaluator. It has 2 tensors31// of shape {1}, and 'evaluation' adds them.32// The interface is the one expected by ReleaseModelRunner.33class MockAOTModelBase {34protected:35 int64_t A = 0;36 int64_t B = 0;37 int64_t R = 0;38 39public:40 MockAOTModelBase() = default;41 virtual ~MockAOTModelBase() = default;42 43 virtual int LookupArgIndex(const std::string &Name) {44 if (Name == "prefix_a")45 return 0;46 if (Name == "prefix_b")47 return 1;48 return -1;49 }50 int LookupResultIndex(const std::string &) { return 0; }51 virtual void Run() = 0;52 virtual void *result_data(int RIndex) {53 if (RIndex == 0)54 return &R;55 return nullptr;56 }57 virtual void *arg_data(int Index) {58 switch (Index) {59 case 0:60 return &A;61 case 1:62 return &B;63 default:64 return nullptr;65 }66 }67};68 69class AdditionAOTModel final : public MockAOTModelBase {70public:71 AdditionAOTModel() = default;72 void Run() override { R = A + B; }73};74 75class DiffAOTModel final : public MockAOTModelBase {76public:77 DiffAOTModel() = default;78 void Run() override { R = A - B; }79};80 81static const char *M1Selector = "the model that subtracts";82static const char *M2Selector = "the model that adds";83 84static MD5::MD5Result Hash1 = MD5::hash(arrayRefFromStringRef(M1Selector));85static MD5::MD5Result Hash2 = MD5::hash(arrayRefFromStringRef(M2Selector));86class ComposedAOTModel final {87 DiffAOTModel M1;88 AdditionAOTModel M2;89 uint64_t Selector[2] = {0};90 91 bool isHashSameAsSelector(const std::pair<uint64_t, uint64_t> &Words) const {92 return Selector[0] == Words.first && Selector[1] == Words.second;93 }94 MockAOTModelBase *getModel() {95 if (isHashSameAsSelector(Hash1.words()))96 return &M1;97 if (isHashSameAsSelector(Hash2.words()))98 return &M2;99 llvm_unreachable("Should be one of the two");100 }101 102public:103 ComposedAOTModel() = default;104 int LookupArgIndex(const std::string &Name) {105 if (Name == "prefix_model_selector")106 return 2;107 return getModel()->LookupArgIndex(Name);108 }109 int LookupResultIndex(const std::string &Name) {110 return getModel()->LookupResultIndex(Name);111 }112 void *arg_data(int Index) {113 if (Index == 2)114 return Selector;115 return getModel()->arg_data(Index);116 }117 void *result_data(int RIndex) { return getModel()->result_data(RIndex); }118 void Run() { getModel()->Run(); }119};120 121static EmbeddedModelRunnerOptions makeOptions() {122 EmbeddedModelRunnerOptions Opts;123 Opts.setFeedPrefix("prefix_");124 return Opts;125}126} // namespace llvm127 128TEST(NoInferenceModelRunner, AccessTensors) {129 const std::vector<TensorSpec> Inputs{130 TensorSpec::createSpec<int64_t>("F1", {1}),131 TensorSpec::createSpec<int64_t>("F2", {10}),132 TensorSpec::createSpec<float>("F2", {5}),133 };134 LLVMContext Ctx;135 NoInferenceModelRunner NIMR(Ctx, Inputs);136 NIMR.getTensor<int64_t>(0)[0] = 1;137 std::memcpy(NIMR.getTensor<int64_t>(1),138 std::vector<int64_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(),139 10 * sizeof(int64_t));140 std::memcpy(NIMR.getTensor<float>(2),141 std::vector<float>{0.1f, 0.2f, 0.3f, 0.4f, 0.5f}.data(),142 5 * sizeof(float));143 ASSERT_EQ(NIMR.getTensor<int64_t>(0)[0], 1);144 ASSERT_EQ(NIMR.getTensor<int64_t>(1)[8], 9);145 ASSERT_EQ(NIMR.getTensor<float>(2)[1], 0.2f);146}147 148TEST(ReleaseModeRunner, NormalUse) {149 LLVMContext Ctx;150 std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),151 TensorSpec::createSpec<int64_t>("b", {1})};152 auto Evaluator = std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(153 Ctx, Inputs, "", makeOptions());154 *Evaluator->getTensor<int64_t>(0) = 1;155 *Evaluator->getTensor<int64_t>(1) = 2;156 EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);157 EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);158 EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);159}160 161TEST(ReleaseModeRunner, ExtraFeatures) {162 LLVMContext Ctx;163 std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),164 TensorSpec::createSpec<int64_t>("b", {1}),165 TensorSpec::createSpec<int64_t>("c", {1})};166 auto Evaluator = std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(167 Ctx, Inputs, "", makeOptions());168 *Evaluator->getTensor<int64_t>(0) = 1;169 *Evaluator->getTensor<int64_t>(1) = 2;170 *Evaluator->getTensor<int64_t>(2) = -3;171 EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);172 EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);173 EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);174 EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);175}176 177TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {178 LLVMContext Ctx;179 std::vector<TensorSpec> Inputs{180 TensorSpec::createSpec<int64_t>("a", {1}),181 TensorSpec::createSpec<int64_t>("c", {1}),182 TensorSpec::createSpec<int64_t>("b", {1}),183 };184 auto Evaluator = std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(185 Ctx, Inputs, "", makeOptions());186 *Evaluator->getTensor<int64_t>(0) = 1; // a187 *Evaluator->getTensor<int64_t>(1) = 2; // c188 *Evaluator->getTensor<int64_t>(2) = -3; // b189 EXPECT_EQ(Evaluator->evaluate<int64_t>(), -2); // a + b190 EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);191 EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);192 EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);193}194 195// We expect an error to be reported early if the user tried to specify a model196// selector, but the model in fact doesn't support that.197TEST(ReleaseModelRunner, ModelSelectorNoInputFeaturePresent) {198 LLVMContext Ctx;199 std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),200 TensorSpec::createSpec<int64_t>("b", {1})};201 EXPECT_DEATH((void)std::make_unique<ReleaseModeModelRunner<AdditionAOTModel>>(202 Ctx, Inputs, "", makeOptions().setModelSelector(M2Selector)),203 "A model selector was specified but the underlying model does "204 "not expose a model_selector input");205}206 207TEST(ReleaseModelRunner, ModelSelectorNoSelectorGiven) {208 LLVMContext Ctx;209 std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),210 TensorSpec::createSpec<int64_t>("b", {1})};211 EXPECT_DEATH(212 (void)std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(213 Ctx, Inputs, "", makeOptions()),214 "A model selector was not specified but the underlying model requires "215 "selecting one because it exposes a model_selector input");216}217 218// Test that we correctly set up the model_selector tensor value. We are only219// responsbile for what happens if the user doesn't specify a value (but the220// model supports the feature), or if the user specifies one, and we correctly221// populate the tensor, and do so upfront (in case the model implementation222// needs that for subsequent tensor buffer lookups).223TEST(ReleaseModelRunner, ModelSelector) {224 LLVMContext Ctx;225 std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),226 TensorSpec::createSpec<int64_t>("b", {1})};227 // This explicitly asks for M1228 auto Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(229 Ctx, Inputs, "", makeOptions().setModelSelector(M1Selector));230 *Evaluator->getTensor<int64_t>(0) = 1;231 *Evaluator->getTensor<int64_t>(1) = 2;232 EXPECT_EQ(Evaluator->evaluate<int64_t>(), -1);233 234 // Ask for M2235 Evaluator = std::make_unique<ReleaseModeModelRunner<ComposedAOTModel>>(236 Ctx, Inputs, "", makeOptions().setModelSelector(M2Selector));237 *Evaluator->getTensor<int64_t>(0) = 1;238 *Evaluator->getTensor<int64_t>(1) = 2;239 EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);240 241 // Asking for a model that's not supported isn't handled by our infra and we242 // expect the model implementation to fail at a point.243}244 245#if defined(LLVM_ON_UNIX)246TEST(InteractiveModelRunner, Evaluation) {247 LLVMContext Ctx;248 // Test the interaction with an external advisor by asking for advice twice.249 // Use simple values, since we use the Logger underneath, that's tested more250 // extensively elsewhere.251 std::vector<TensorSpec> Inputs{252 TensorSpec::createSpec<int64_t>("a", {1}),253 TensorSpec::createSpec<int64_t>("b", {1}),254 TensorSpec::createSpec<int64_t>("c", {1}),255 };256 TensorSpec AdviceSpec = TensorSpec::createSpec<float>("advice", {1});257 258 // Create the 2 files. Ideally we'd create them as named pipes, but that's not259 // quite supported by the generic API.260 std::error_code EC;261 llvm::unittest::TempDir Tmp("tmpdir", /*Unique=*/true);262 SmallString<128> FromCompilerName(Tmp.path().begin(), Tmp.path().end());263 SmallString<128> ToCompilerName(Tmp.path().begin(), Tmp.path().end());264 sys::path::append(FromCompilerName, "InteractiveModelRunner_Evaluation.out");265 sys::path::append(ToCompilerName, "InteractiveModelRunner_Evaluation.in");266 EXPECT_EQ(::mkfifo(FromCompilerName.c_str(), 0666), 0);267 EXPECT_EQ(::mkfifo(ToCompilerName.c_str(), 0666), 0);268 269 FileRemover Cleanup1(FromCompilerName);270 FileRemover Cleanup2(ToCompilerName);271 272 // Since the evaluator sends the features over and then blocks waiting for273 // an answer, we must spawn a thread playing the role of the advisor / host:274 std::atomic<int> SeenObservations = 0;275 // Start the host first to make sure the pipes are being prepared. Otherwise276 // the evaluator will hang.277 std::thread Advisor([&]() {278 // Open the writer first. This is because the evaluator will try opening279 // the "input" pipe first. An alternative that avoids ordering is for the280 // host to open the pipes RW.281 raw_fd_ostream ToCompiler(ToCompilerName, EC);282 EXPECT_FALSE(EC);283 int FromCompilerHandle = 0;284 EXPECT_FALSE(285 sys::fs::openFileForRead(FromCompilerName, FromCompilerHandle));286 sys::fs::file_t FromCompiler =287 sys::fs::convertFDToNativeFile(FromCompilerHandle);288 EXPECT_EQ(SeenObservations, 0);289 // Helper to read headers and other json lines.290 SmallVector<char, 1024> Buffer;291 auto ReadLn = [&]() {292 Buffer.clear();293 while (true) {294 char Chr = 0;295 auto ReadOrErr = sys::fs::readNativeFile(FromCompiler, {&Chr, 1});296 EXPECT_FALSE(ReadOrErr.takeError());297 if (!*ReadOrErr)298 continue;299 if (Chr == '\n')300 return StringRef(Buffer.data(), Buffer.size());301 Buffer.push_back(Chr);302 }303 };304 // See include/llvm/Analysis/Utils/TrainingLogger.h305 // First comes the header306 auto Header = json::parse(ReadLn());307 EXPECT_FALSE(Header.takeError());308 EXPECT_NE(Header->getAsObject()->getArray("features"), nullptr);309 EXPECT_NE(Header->getAsObject()->getObject("advice"), nullptr);310 // Then comes the context311 EXPECT_FALSE(json::parse(ReadLn()).takeError());312 313 int64_t Features[3] = {0};314 auto FullyRead = [&]() {315 size_t InsPt = 0;316 const size_t ToRead = 3 * Inputs[0].getTotalTensorBufferSize();317 char *Buff = reinterpret_cast<char *>(Features);318 while (InsPt < ToRead) {319 auto ReadOrErr = sys::fs::readNativeFile(320 FromCompiler, {Buff + InsPt, ToRead - InsPt});321 EXPECT_FALSE(ReadOrErr.takeError());322 InsPt += *ReadOrErr;323 }324 };325 // Observation326 EXPECT_FALSE(json::parse(ReadLn()).takeError());327 // Tensor values328 FullyRead();329 // a "\n"330 char Chr = 0;331 auto ReadNL = [&]() {332 do {333 auto ReadOrErr = sys::fs::readNativeFile(FromCompiler, {&Chr, 1});334 EXPECT_FALSE(ReadOrErr.takeError());335 if (*ReadOrErr == 1)336 break;337 } while (true);338 };339 ReadNL();340 EXPECT_EQ(Chr, '\n');341 EXPECT_EQ(Features[0], 42);342 EXPECT_EQ(Features[1], 43);343 EXPECT_EQ(Features[2], 100);344 ++SeenObservations;345 346 // Send the advice347 float Advice = 42.0012;348 ToCompiler.write(reinterpret_cast<const char *>(&Advice),349 AdviceSpec.getTotalTensorBufferSize());350 ToCompiler.flush();351 352 // Second observation, and same idea as above353 EXPECT_FALSE(json::parse(ReadLn()).takeError());354 FullyRead();355 ReadNL();356 EXPECT_EQ(Chr, '\n');357 EXPECT_EQ(Features[0], 10);358 EXPECT_EQ(Features[1], -2);359 EXPECT_EQ(Features[2], 1);360 ++SeenObservations;361 Advice = 50.30;362 ToCompiler.write(reinterpret_cast<const char *>(&Advice),363 AdviceSpec.getTotalTensorBufferSize());364 ToCompiler.flush();365 sys::fs::closeFile(FromCompiler);366 });367 368 InteractiveModelRunner Evaluator(Ctx, Inputs, AdviceSpec, FromCompilerName,369 ToCompilerName);370 371 Evaluator.switchContext("hi");372 373 EXPECT_EQ(SeenObservations, 0);374 *Evaluator.getTensor<int64_t>(0) = 42;375 *Evaluator.getTensor<int64_t>(1) = 43;376 *Evaluator.getTensor<int64_t>(2) = 100;377 float Ret = Evaluator.evaluate<float>();378 EXPECT_EQ(SeenObservations, 1);379 EXPECT_FLOAT_EQ(Ret, 42.0012);380 381 *Evaluator.getTensor<int64_t>(0) = 10;382 *Evaluator.getTensor<int64_t>(1) = -2;383 *Evaluator.getTensor<int64_t>(2) = 1;384 Ret = Evaluator.evaluate<float>();385 EXPECT_EQ(SeenObservations, 2);386 EXPECT_FLOAT_EQ(Ret, 50.30);387 Advisor.join();388}389#endif390