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>(Name: "F1" , Shape: {1}), |
68 | TensorSpec::createSpec<int64_t>(Name: "F2" , Shape: {10}), |
69 | TensorSpec::createSpec<float>(Name: "F2" , Shape: {5}), |
70 | }; |
71 | LLVMContext Ctx; |
72 | NoInferenceModelRunner NIMR(Ctx, Inputs); |
73 | NIMR.getTensor<int64_t>(FeatureID: 0)[0] = 1; |
74 | std::memcpy(dest: NIMR.getTensor<int64_t>(FeatureID: 1), |
75 | src: std::vector<int64_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(), |
76 | n: 10 * sizeof(int64_t)); |
77 | std::memcpy(dest: NIMR.getTensor<float>(FeatureID: 2), |
78 | src: std::vector<float>{0.1f, 0.2f, 0.3f, 0.4f, 0.5f}.data(), |
79 | n: 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>(Name: "a" , Shape: {1}), |
88 | TensorSpec::createSpec<int64_t>(Name: "b" , Shape: {1})}; |
89 | auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>( |
90 | args&: Ctx, args&: Inputs, args: "" , args: "prefix_" ); |
91 | *Evaluator->getTensor<int64_t>(FeatureID: 0) = 1; |
92 | *Evaluator->getTensor<int64_t>(FeatureID: 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>(Name: "a" , Shape: {1}), |
101 | TensorSpec::createSpec<int64_t>(Name: "b" , Shape: {1}), |
102 | TensorSpec::createSpec<int64_t>(Name: "c" , Shape: {1})}; |
103 | auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>( |
104 | args&: Ctx, args&: Inputs, args: "" , args: "prefix_" ); |
105 | *Evaluator->getTensor<int64_t>(FeatureID: 0) = 1; |
106 | *Evaluator->getTensor<int64_t>(FeatureID: 1) = 2; |
107 | *Evaluator->getTensor<int64_t>(FeatureID: 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>(Name: "a" , Shape: {1}), |
118 | TensorSpec::createSpec<int64_t>(Name: "c" , Shape: {1}), |
119 | TensorSpec::createSpec<int64_t>(Name: "b" , Shape: {1}), |
120 | }; |
121 | auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>( |
122 | args&: Ctx, args&: Inputs, args: "" , args: "prefix_" ); |
123 | *Evaluator->getTensor<int64_t>(FeatureID: 0) = 1; // a |
124 | *Evaluator->getTensor<int64_t>(FeatureID: 1) = 2; // c |
125 | *Evaluator->getTensor<int64_t>(FeatureID: 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>(Name: "a" , Shape: {1}), |
140 | TensorSpec::createSpec<int64_t>(Name: "b" , Shape: {1}), |
141 | TensorSpec::createSpec<int64_t>(Name: "c" , Shape: {1}), |
142 | }; |
143 | TensorSpec AdviceSpec = TensorSpec::createSpec<float>(Name: "advice" , Shape: {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(path&: FromCompilerName, a: "InteractiveModelRunner_Evaluation.out" ); |
152 | sys::path::append(path&: ToCompilerName, a: "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(FD: 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(FileHandle: FromCompiler, Buf: {&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(Elt: Chr); |
189 | } |
190 | }; |
191 | // See include/llvm/Analysis/Utils/TrainingLogger.h |
192 | // First comes the header |
193 | auto = json::parse(JSON: 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 | FileHandle: FromCompiler, Buf: {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(FileHandle: FromCompiler, Buf: {&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(Ptr: reinterpret_cast<const char *>(&Advice), |
236 | Size: 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(Ptr: reinterpret_cast<const char *>(&Advice), |
250 | Size: AdviceSpec.getTotalTensorBufferSize()); |
251 | ToCompiler.flush(); |
252 | sys::fs::closeFile(F&: FromCompiler); |
253 | }); |
254 | |
255 | InteractiveModelRunner Evaluator(Ctx, Inputs, AdviceSpec, FromCompilerName, |
256 | ToCompilerName); |
257 | |
258 | Evaluator.switchContext(Name: "hi" ); |
259 | |
260 | EXPECT_EQ(SeenObservations, 0); |
261 | *Evaluator.getTensor<int64_t>(FeatureID: 0) = 42; |
262 | *Evaluator.getTensor<int64_t>(FeatureID: 1) = 43; |
263 | *Evaluator.getTensor<int64_t>(FeatureID: 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>(FeatureID: 0) = 10; |
269 | *Evaluator.getTensor<int64_t>(FeatureID: 1) = -2; |
270 | *Evaluator.getTensor<int64_t>(FeatureID: 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 |
277 | |