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