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/Analysis/InteractiveModelRunner.h"
  11 #include "llvm/Analysis/NoInferenceModelRunner.h"
  12 #include "llvm/Analysis/ReleaseModeModelRunner.h"
  13 #include "llvm/Support/BinaryByteStream.h"
  14 #include "llvm/Support/FileSystem.h"
  15 #include "llvm/Support/FileUtilities.h"
  16 #include "llvm/Support/JSON.h"
  17 #include "llvm/Support/Path.h"
  18 #include "llvm/Support/raw_ostream.h"
  19 #include "llvm/Testing/Support/SupportHelpers.h"
  20 #include "gtest/gtest.h"
  21
  22 #include <atomic>
  23 #include <thread>
  24
  25 using namespace llvm;
  26
  27 namespace llvm {
  28 // This is a mock of the kind of AOT-generated model evaluator. It has 2 tensors
  29 // of shape {1}, and 'evaluation' adds them.
  30 // The interface is the one expected by ReleaseModelRunner.
  31 class MockAOTModel final {
  32   int64_t A = 0;
  33   int64_t B = 0;
  34   int64_t R = 0;
  35
  36 public:
  37   MockAOTModel() = default;
  38   int LookupArgIndex(const std::string &Name) {
  39     if (Name == "prefix_a")
  40       return 0;
  41     if (Name == "prefix_b")
  42       return 1;
  43     return -1;
  44   }
  45   int LookupResultIndex(const std::string &) { return 0; }
  46   void Run() { R = A + B; }
  47   void *result_data(int RIndex) {
  48     if (RIndex == 0)
  49       return &R;
  50     return nullptr;
  51   }
  52   void *arg_data(int Index) {
  53     switch (Index) {
  54     case 0:
  55       return &A;
  56     case 1:
  57       return &B;
  58     default:
  59       return nullptr;
  60     }
  61   }
  62 };
  63 } // namespace llvm
  64
  65 TEST(NoInferenceModelRunner, AccessTensors) {
  66   const std::vector<TensorSpec> Inputs{
  67       TensorSpec::createSpec<int64_t>("F1", {1}),
  68       TensorSpec::createSpec<int64_t>("F2", {10}),
  69       TensorSpec::createSpec<float>("F2", {5}),
  70   };
  71   LLVMContext Ctx;
  72   NoInferenceModelRunner NIMR(Ctx, Inputs);
  73   NIMR.getTensor<int64_t>(0)[0] = 1;
  74   std::memcpy(NIMR.getTensor<int64_t>(1),
  75               std::vector<int64_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(),
  76               10 * sizeof(int64_t));
  77   std::memcpy(NIMR.getTensor<float>(2),
  78               std::vector<float>{0.1f, 0.2f, 0.3f, 0.4f, 0.5f}.data(),
  79               5 * sizeof(float));
  80   ASSERT_EQ(NIMR.getTensor<int64_t>(0)[0], 1);
  81   ASSERT_EQ(NIMR.getTensor<int64_t>(1)[8], 9);
  82   ASSERT_EQ(NIMR.getTensor<float>(2)[1], 0.2f);
  83 }
  84
  85 TEST(ReleaseModeRunner, NormalUse) {
  86   LLVMContext Ctx;
  87   std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
  88                                  TensorSpec::createSpec<int64_t>("b", {1})};
  89   auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
  90       Ctx, Inputs, "", "prefix_");
  91   *Evaluator->getTensor<int64_t>(0) = 1;
  92   *Evaluator->getTensor<int64_t>(1) = 2;
  93   EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
  94   EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
  95   EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
  96 }
  97
  98 TEST(ReleaseModeRunner, ExtraFeatures) {
  99   LLVMContext Ctx;
 100   std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
 101                                  TensorSpec::createSpec<int64_t>("b", {1}),
 102                                  TensorSpec::createSpec<int64_t>("c", {1})};
 103   auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
 104       Ctx, Inputs, "", "prefix_");
 105   *Evaluator->getTensor<int64_t>(0) = 1;
 106   *Evaluator->getTensor<int64_t>(1) = 2;
 107   *Evaluator->getTensor<int64_t>(2) = -3;
 108   EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
 109   EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
 110   EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
 111   EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
 112 }
 113
 114 TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {
 115   LLVMContext Ctx;
 116   std::vector<TensorSpec> Inputs{
 117       TensorSpec::createSpec<int64_t>("a", {1}),
 118       TensorSpec::createSpec<int64_t>("c", {1}),
 119       TensorSpec::createSpec<int64_t>("b", {1}),
 120   };
 121   auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
 122       Ctx, Inputs, "", "prefix_");
 123   *Evaluator->getTensor<int64_t>(0) = 1;         // a
 124   *Evaluator->getTensor<int64_t>(1) = 2;         // c
 125   *Evaluator->getTensor<int64_t>(2) = -3;        // b
 126   EXPECT_EQ(Evaluator->evaluate<int64_t>(), -2); // a + b
 127   EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
 128   EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
 129   EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
 130 }
 131
 132 #if defined(LLVM_ON_UNIX)
 133 TEST(InteractiveModelRunner, Evaluation) {
 134   LLVMContext Ctx;
 135   // Test the interaction with an external advisor by asking for advice twice.
 136   // Use simple values, since we use the Logger underneath, that's tested more
 137   // extensively elsewhere.
 138   std::vector<TensorSpec> Inputs{
 139       TensorSpec::createSpec<int64_t>("a", {1}),
 140       TensorSpec::createSpec<int64_t>("b", {1}),
 141       TensorSpec::createSpec<int64_t>("c", {1}),
 142   };
 143   TensorSpec AdviceSpec = TensorSpec::createSpec<float>("advice", {1});
 144
 145   // Create the 2 files. Ideally we'd create them as named pipes, but that's not
 146   // quite supported by the generic API.
 147   std::error_code EC;
 148   llvm::unittest::TempDir Tmp("tmpdir", /*Unique=*/true);
 149   SmallString<128> FromCompilerName(Tmp.path().begin(), Tmp.path().end());
 150   SmallString<128> ToCompilerName(Tmp.path().begin(), Tmp.path().end());
 151   sys::path::append(FromCompilerName, "InteractiveModelRunner_Evaluation.out");
 152   sys::path::append(ToCompilerName, "InteractiveModelRunner_Evaluation.in");
 153   EXPECT_EQ(::mkfifo(FromCompilerName.c_str(), 0666), 0);
 154   EXPECT_EQ(::mkfifo(ToCompilerName.c_str(), 0666), 0);
 155
 156   FileRemover Cleanup1(FromCompilerName);
 157   FileRemover Cleanup2(ToCompilerName);
 158
 159   // Since the evaluator sends the features over and then blocks waiting for
 160   // an answer, we must spawn a thread playing the role of the advisor / host:
 161   std::atomic<int> SeenObservations = 0;
 162   // Start the host first to make sure the pipes are being prepared. Otherwise
 163   // the evaluator will hang.
 164   std::thread Advisor([&]() {
 165     // Open the writer first. This is because the evaluator will try opening
 166     // the "input" pipe first. An alternative that avoids ordering is for the
 167     // host to open the pipes RW.
 168     raw_fd_ostream ToCompiler(ToCompilerName, EC);
 169     EXPECT_FALSE(EC);
 170     int FromCompilerHandle = 0;
 171     EXPECT_FALSE(
 172         sys::fs::openFileForRead(FromCompilerName, FromCompilerHandle));
 173     sys::fs::file_t FromCompiler =
 174         sys::fs::convertFDToNativeFile(FromCompilerHandle);
 175     EXPECT_EQ(SeenObservations, 0);
 176     // Helper to read headers and other json lines.
 177     SmallVector<char, 1024> Buffer;
 178     auto ReadLn = [&]() {
 179       Buffer.clear();
 180       while (true) {
 181         char Chr = 0;
 182         auto ReadOrErr = sys::fs::readNativeFile(FromCompiler, {&Chr, 1});
 183         EXPECT_FALSE(ReadOrErr.takeError());
 184         if (!*ReadOrErr)
 185           continue;
 186         if (Chr == '\n')
 187           return StringRef(Buffer.data(), Buffer.size());
 188         Buffer.push_back(Chr);
 189       }
 190     };
 191     // See include/llvm/Analysis/Utils/TrainingLogger.h
 192     // First comes the header
 193     auto Header = json::parse(ReadLn());
 194     EXPECT_FALSE(Header.takeError());
 195     EXPECT_NE(Header->getAsObject()->getArray("features"), nullptr);
 196     EXPECT_NE(Header->getAsObject()->getObject("advice"), nullptr);
 197     // Then comes the context
 198     EXPECT_FALSE(json::parse(ReadLn()).takeError());
 199
 200     int64_t Features[3] = {0};
 201     auto FullyRead = [&]() {
 202       size_t InsPt = 0;
 203       const size_t ToRead = 3 * Inputs[0].getTotalTensorBufferSize();
 204       char *Buff = reinterpret_cast<char *>(Features);
 205       while (InsPt < ToRead) {
 206         auto ReadOrErr = sys::fs::readNativeFile(
 207             FromCompiler, {Buff + InsPt, ToRead - InsPt});
 208         EXPECT_FALSE(ReadOrErr.takeError());
 209         InsPt += *ReadOrErr;
 210       }
 211     };
 212     // Observation
 213     EXPECT_FALSE(json::parse(ReadLn()).takeError());
 214     // Tensor values
 215     FullyRead();
 216     // a "\n"
 217     char Chr = 0;
 218     auto ReadNL = [&]() {
 219       do {
 220         auto ReadOrErr = sys::fs::readNativeFile(FromCompiler, {&Chr, 1});
 221         EXPECT_FALSE(ReadOrErr.takeError());
 222         if (*ReadOrErr == 1)
 223           break;
 224       } while (true);
 225     };
 226     ReadNL();
 227     EXPECT_EQ(Chr, '\n');
 228     EXPECT_EQ(Features[0], 42);
 229     EXPECT_EQ(Features[1], 43);
 230     EXPECT_EQ(Features[2], 100);
 231     ++SeenObservations;
 232
 233     // Send the advice
 234     float Advice = 42.0012;
 235     ToCompiler.write(reinterpret_cast<const char *>(&Advice),
 236                      AdviceSpec.getTotalTensorBufferSize());
 237     ToCompiler.flush();
 238
 239     // Second observation, and same idea as above
 240     EXPECT_FALSE(json::parse(ReadLn()).takeError());
 241     FullyRead();
 242     ReadNL();
 243     EXPECT_EQ(Chr, '\n');
 244     EXPECT_EQ(Features[0], 10);
 245     EXPECT_EQ(Features[1], -2);
 246     EXPECT_EQ(Features[2], 1);
 247     ++SeenObservations;
 248     Advice = 50.30;
 249     ToCompiler.write(reinterpret_cast<const char *>(&Advice),
 250                      AdviceSpec.getTotalTensorBufferSize());
 251     ToCompiler.flush();
 252     sys::fs::closeFile(FromCompiler);
 253   });
 254
 255   InteractiveModelRunner Evaluator(Ctx, Inputs, AdviceSpec, FromCompilerName,
 256                                    ToCompilerName);
 257
 258   Evaluator.switchContext("hi");
 259
 260   EXPECT_EQ(SeenObservations, 0);
 261   *Evaluator.getTensor<int64_t>(0) = 42;
 262   *Evaluator.getTensor<int64_t>(1) = 43;
 263   *Evaluator.getTensor<int64_t>(2) = 100;
 264   float Ret = Evaluator.evaluate<float>();
 265   EXPECT_EQ(SeenObservations, 1);
 266   EXPECT_FLOAT_EQ(Ret, 42.0012);
 267
 268   *Evaluator.getTensor<int64_t>(0) = 10;
 269   *Evaluator.getTensor<int64_t>(1) = -2;
 270   *Evaluator.getTensor<int64_t>(2) = 1;
 271   Ret = Evaluator.evaluate<float>();
 272   EXPECT_EQ(SeenObservations, 2);
 273   EXPECT_FLOAT_EQ(Ret, 50.30);
 274   Advisor.join();
 275 }
 276 #endif