llvm/unittests/Analysis/MLModelRunnerTest.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-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