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 sys::fs::file_t FromCompiler = {}; 171 EXPECT_FALSE(sys::fs::openFileForRead(FromCompilerName, FromCompiler)); 172 EXPECT_EQ(SeenObservations, 0); 173 // Helper to read headers and other json lines. 174 SmallVector<char, 1024> Buffer; 175 auto ReadLn = [&]() { 176 Buffer.clear(); 177 while (true) { 178 char Chr = 0; 179 auto ReadOrErr = sys::fs::readNativeFile(FromCompiler, {&Chr, 1}); 180 EXPECT_FALSE(ReadOrErr.takeError()); 181 if (!*ReadOrErr) 182 continue; 183 if (Chr == '\n') 184 return StringRef(Buffer.data(), Buffer.size()); 185 Buffer.push_back(Chr); 186 } 187 }; 188 // See include/llvm/Analysis/Utils/TrainingLogger.h 189 // First comes the header 190 auto Header = json::parse(ReadLn()); 191 EXPECT_FALSE(Header.takeError()); 192 EXPECT_NE(Header->getAsObject()->getArray("features"), nullptr); 193 EXPECT_NE(Header->getAsObject()->getObject("advice"), nullptr); 194 // Then comes the context 195 EXPECT_FALSE(json::parse(ReadLn()).takeError()); 196 197 int64_t Features[3] = {0}; 198 auto FullyRead = [&]() { 199 size_t InsPt = 0; 200 const size_t ToRead = 3 * Inputs[0].getTotalTensorBufferSize(); 201 char *Buff = reinterpret_cast<char *>(Features); 202 while (InsPt < ToRead) { 203 auto ReadOrErr = sys::fs::readNativeFile( 204 FromCompiler, {Buff + InsPt, ToRead - InsPt}); 205 EXPECT_FALSE(ReadOrErr.takeError()); 206 InsPt += *ReadOrErr; 207 } 208 }; 209 // Observation 210 EXPECT_FALSE(json::parse(ReadLn()).takeError()); 211 // Tensor values 212 FullyRead(); 213 // a "\n" 214 char Chr = 0; 215 auto ReadNL = [&]() { 216 do { 217 auto ReadOrErr = sys::fs::readNativeFile(FromCompiler, {&Chr, 1}); 218 EXPECT_FALSE(ReadOrErr.takeError()); 219 if (*ReadOrErr == 1) 220 break; 221 } while (true); 222 }; 223 ReadNL(); 224 EXPECT_EQ(Chr, '\n'); 225 EXPECT_EQ(Features[0], 42); 226 EXPECT_EQ(Features[1], 43); 227 EXPECT_EQ(Features[2], 100); 228 ++SeenObservations; 229 230 // Send the advice 231 float Advice = 42.0012; 232 ToCompiler.write(reinterpret_cast<const char *>(&Advice), 233 AdviceSpec.getTotalTensorBufferSize()); 234 ToCompiler.flush(); 235 236 // Second observation, and same idea as above 237 EXPECT_FALSE(json::parse(ReadLn()).takeError()); 238 FullyRead(); 239 ReadNL(); 240 EXPECT_EQ(Chr, '\n'); 241 EXPECT_EQ(Features[0], 10); 242 EXPECT_EQ(Features[1], -2); 243 EXPECT_EQ(Features[2], 1); 244 ++SeenObservations; 245 Advice = 50.30; 246 ToCompiler.write(reinterpret_cast<const char *>(&Advice), 247 AdviceSpec.getTotalTensorBufferSize()); 248 ToCompiler.flush(); 249 sys::fs::closeFile(FromCompiler); 250 }); 251 252 InteractiveModelRunner Evaluator(Ctx, Inputs, AdviceSpec, FromCompilerName, 253 ToCompilerName); 254 255 Evaluator.switchContext("hi"); 256 257 EXPECT_EQ(SeenObservations, 0); 258 *Evaluator.getTensor<int64_t>(0) = 42; 259 *Evaluator.getTensor<int64_t>(1) = 43; 260 *Evaluator.getTensor<int64_t>(2) = 100; 261 float Ret = Evaluator.evaluate<float>(); 262 EXPECT_EQ(SeenObservations, 1); 263 EXPECT_FLOAT_EQ(Ret, 42.0012); 264 265 *Evaluator.getTensor<int64_t>(0) = 10; 266 *Evaluator.getTensor<int64_t>(1) = -2; 267 *Evaluator.getTensor<int64_t>(2) = 1; 268 Ret = Evaluator.evaluate<float>(); 269 EXPECT_EQ(SeenObservations, 2); 270 EXPECT_FLOAT_EQ(Ret, 50.30); 271 Advisor.join(); 272 } 273 #endif 274