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-exception 6 // 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_UNIX 15 #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 27 using namespace llvm; 28 29 namespace llvm { 30 // This is a mock of the kind of AOT-generated model evaluator. It has 2 tensors 31 // of shape {1}, and 'evaluation' adds them. 32 // The interface is the one expected by ReleaseModelRunner. 33 class MockAOTModelBase { 34 protected: 35 int64_t A = 0; 36 int64_t B = 0; 37 int64_t R = 0; 38 39 public: 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 69 class AdditionAOTModel final : public MockAOTModelBase { 70 public: 71 AdditionAOTModel() = default; 72 void Run() override { R = A + B; } 73 }; 74 75 class DiffAOTModel final : public MockAOTModelBase { 76 public: 77 DiffAOTModel() = default; 78 void Run() override { R = A - B; } 79 }; 80 81 static const char *M1Selector = "the model that subtracts"; 82 static const char *M2Selector = "the model that adds"; 83 84 static MD5::MD5Result Hash1 = MD5::hash(arrayRefFromStringRef(M1Selector)); 85 static MD5::MD5Result Hash2 = MD5::hash(arrayRefFromStringRef(M2Selector)); 86 class 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 102 public: 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 121 static EmbeddedModelRunnerOptions makeOptions() { 122 EmbeddedModelRunnerOptions Opts; 123 Opts.setFeedPrefix("prefix_"); 124 return Opts; 125 } 126 } // namespace llvm 127 128 TEST(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 148 TEST(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 161 TEST(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 177 TEST(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; // a 187 *Evaluator->getTensor<int64_t>(1) = 2; // c 188 *Evaluator->getTensor<int64_t>(2) = -3; // b 189 EXPECT_EQ(Evaluator->evaluate<int64_t>(), -2); // a + b 190 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 model 196 // selector, but the model in fact doesn't support that. 197 TEST(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 207 TEST(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 only 219 // responsbile for what happens if the user doesn't specify a value (but the 220 // model supports the feature), or if the user specifies one, and we correctly 221 // populate the tensor, and do so upfront (in case the model implementation 222 // needs that for subsequent tensor buffer lookups). 223 TEST(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 M1 228 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 M2 235 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 we 242 // expect the model implementation to fail at a point. 243 } 244 245 #if defined(LLVM_ON_UNIX) 246 TEST(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 more 250 // 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 not 259 // 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 for 273 // 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. Otherwise 276 // the evaluator will hang. 277 std::thread Advisor([&]() { 278 // Open the writer first. This is because the evaluator will try opening 279 // the "input" pipe first. An alternative that avoids ordering is for the 280 // 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.h 305 // First comes the header 306 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 context 311 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 // Observation 326 EXPECT_FALSE(json::parse(ReadLn()).takeError()); 327 // Tensor values 328 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 advice 347 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 above 353 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 #endif 390