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
25using namespace llvm;
26
27namespace 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.
31class MockAOTModel final {
32 int64_t A = 0;
33 int64_t B = 0;
34 int64_t R = 0;
35
36public:
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
65TEST(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
85TEST(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
98TEST(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
114TEST(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)
133TEST(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 Header = 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

source code of llvm/unittests/Analysis/MLModelRunnerTest.cpp