1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <cstdlib>
10 #include <filesystem>
11
12 #include <executorch/extension/data_loader/file_data_loader.h>
13 #include <executorch/extension/runner_util/inputs.h>
14 #include <executorch/runtime/core/exec_aten/exec_aten.h>
15 #include <executorch/runtime/executor/method.h>
16 #include <executorch/runtime/executor/program.h>
17 #include <executorch/runtime/executor/test/managed_memory_manager.h>
18 #include <executorch/runtime/platform/runtime.h>
19 #include <executorch/test/utils/DeathTest.h>
20 #include <gtest/gtest.h>
21
22 using namespace ::testing;
23 using exec_aten::ArrayRef;
24 using executorch::extension::prepare_input_tensors;
25 using executorch::runtime::Error;
26 using executorch::runtime::EValue;
27 using executorch::runtime::Method;
28 using executorch::runtime::Program;
29 using executorch::runtime::Result;
30 using executorch::runtime::testing::ManagedMemoryManager;
31 using torch::executor::util::FileDataLoader;
32
33 constexpr size_t kDefaultNonConstMemBytes = 32 * 1024U;
34 constexpr size_t kDefaultRuntimeMemBytes = 32 * 1024U;
35
36 class MethodTest : public ::testing::Test {
37 protected:
load_program(const char * path,const char * module_name)38 void load_program(const char* path, const char* module_name) {
39 // Create a loader for the serialized program.
40 Result<FileDataLoader> loader = FileDataLoader::from(path);
41 ASSERT_EQ(loader.error(), Error::Ok);
42 loaders_.insert(
43 {module_name,
44 std::make_unique<FileDataLoader>(std::move(loader.get()))});
45
46 // Use it to load the program.
47 Result<Program> program = Program::load(
48 loaders_[module_name].get(),
49 Program::Verification::InternalConsistency);
50 ASSERT_EQ(program.error(), Error::Ok);
51 programs_.insert(
52 {module_name, std::make_unique<Program>(std::move(program.get()))});
53 }
54
SetUp()55 void SetUp() override {
56 executorch::runtime::runtime_init();
57
58 load_program(std::getenv("ET_MODULE_ADD_PATH"), "add");
59 load_program(std::getenv("ET_MODULE_INDEX_PATH"), "index");
60 load_program(
61 std::getenv("ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH"), "cat");
62 load_program(std::getenv("ET_MODULE_LINEAR_PATH"), "linear");
63 load_program(
64 std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH"),
65 "linear_constant_buffer");
66 }
67
68 private:
69 // Must outlive program_, but tests shouldn't need to touch it.
70 std::unordered_map<std::string, std::unique_ptr<FileDataLoader>> loaders_;
71
72 protected:
73 std::unordered_map<std::string, std::unique_ptr<Program>> programs_;
74 };
75
TEST_F(MethodTest,MoveTest)76 TEST_F(MethodTest, MoveTest) {
77 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
78 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
79 ASSERT_EQ(method.error(), Error::Ok);
80
81 // Can execute the method.
82 auto input_cleanup = prepare_input_tensors(*method);
83 ASSERT_EQ(input_cleanup.error(), Error::Ok);
84 Error err = method->execute();
85 ASSERT_EQ(err, Error::Ok);
86
87 // Move into a new Method.
88 Method new_method(std::move(method.get()));
89
90 // Can't execute the old method.
91 err = method->execute();
92 ASSERT_NE(err, Error::Ok);
93
94 // Can execute the new method.
95 err = new_method.execute();
96 ASSERT_EQ(err, Error::Ok);
97 }
98
TEST_F(MethodTest,GetInputTests)99 TEST_F(MethodTest, GetInputTests) {
100 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
101 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
102 ASSERT_EQ(method.error(), Error::Ok);
103
104 size_t num_inputs = method->inputs_size();
105 ASSERT_GT(num_inputs, 0);
106
107 // In-range inputs should succeed without aborting.
108 method->get_input(0);
109 method->get_input(num_inputs - 1);
110
111 // Out-of-range inputs should abort.
112 ET_EXPECT_DEATH(method->get_input(num_inputs), "");
113 ET_EXPECT_DEATH(method->get_input(num_inputs + 1), "");
114 }
115
TEST_F(MethodTest,MutableInputTests)116 TEST_F(MethodTest, MutableInputTests) {
117 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
118 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
119 ASSERT_EQ(method.error(), Error::Ok);
120
121 size_t num_inputs = method->inputs_size();
122 ASSERT_GT(num_inputs, 0);
123
124 // In-range inputs should succeed without aborting.
125 method->mutable_input(0);
126 method->mutable_input(num_inputs - 1);
127
128 // Out-of-range inputs should abort.
129 ET_EXPECT_DEATH(method->mutable_input(num_inputs), "");
130 ET_EXPECT_DEATH(method->mutable_input(num_inputs + 1), "");
131 }
132
TEST_F(MethodTest,GetOutputTests)133 TEST_F(MethodTest, GetOutputTests) {
134 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
135 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
136 ASSERT_EQ(method.error(), Error::Ok);
137
138 size_t num_outputs = method->outputs_size();
139 ASSERT_GT(num_outputs, 0);
140
141 // In-range outputs should succeed without aborting.
142 method->get_output(0);
143 method->get_output(num_outputs - 1);
144
145 // Out-of-range outputs should abort.
146 ET_EXPECT_DEATH(method->get_output(num_outputs), "");
147 ET_EXPECT_DEATH(method->get_output(num_outputs + 1), "");
148 }
149
TEST_F(MethodTest,MutableOutputTests)150 TEST_F(MethodTest, MutableOutputTests) {
151 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
152 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
153 ASSERT_EQ(method.error(), Error::Ok);
154
155 size_t num_outputs = method->outputs_size();
156 ASSERT_GT(num_outputs, 0);
157
158 // In-range outputs should succeed without aborting.
159 method->mutable_output(0);
160 method->mutable_output(num_outputs - 1);
161
162 // Out-of-range outputs should abort.
163 ET_EXPECT_DEATH(method->mutable_output(num_outputs), "");
164 ET_EXPECT_DEATH(method->mutable_output(num_outputs + 1), "");
165 }
166
TEST_F(MethodTest,SetPrimInputTest)167 TEST_F(MethodTest, SetPrimInputTest) {
168 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
169 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
170 ASSERT_EQ(method.error(), Error::Ok);
171
172 // Can execute the method.
173 auto input_cleanup = prepare_input_tensors(*method);
174 ASSERT_EQ(input_cleanup.error(), Error::Ok);
175
176 // The args to the method are x, y, alpha. x and y are tensors handled above
177 // alpha is a prim.
178
179 // Traced prim input was '1.0' so 3.0 should error.
180 auto input_err = method->set_input(EValue(3.0), 2);
181 EXPECT_EQ(input_err, Error::InvalidArgument);
182
183 // Traced prim input was '1.0' so '1.0' should be ok.
184 input_err = method->set_input(EValue(1.0), 2);
185 ASSERT_EQ(input_err, Error::Ok);
186
187 Error err = method->execute();
188 EXPECT_EQ(err, Error::Ok);
189 }
190
TEST_F(MethodTest,MethodMetaTest)191 TEST_F(MethodTest, MethodMetaTest) {
192 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
193 Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
194 ASSERT_EQ(method.error(), Error::Ok);
195
196 auto method_meta = method->method_meta();
197
198 EXPECT_EQ(method_meta.num_inputs(), method->inputs_size());
199 EXPECT_EQ(method_meta.num_outputs(), method->outputs_size());
200 }
201
TEST_F(MethodTest,AliasedIOTest)202 TEST_F(MethodTest, AliasedIOTest) {
203 // TODO(T163238401)
204 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
205 Result<Method> method = programs_["cat"]->load_method("forward", &mmm.get());
206 ASSERT_EQ(method.error(), Error::Ok);
207
208 // Set up io. Input and Output should share the same memory.
209 constexpr int buffer_size = 16;
210 float buffer[buffer_size]; // Initial input is (2,4) we then cat a (1,4) to it
211 // twice for a final shape of (4,4)
212 for (int i = 0; i < buffer_size; ++i) {
213 buffer[i] = 0.f;
214 }
215 int32_t sizes[2] = {2, 4};
216 uint8_t dim_order[2] = {0, 1};
217 int32_t strides[2] = {4, 1};
218 exec_aten::TensorImpl impl(
219 exec_aten::ScalarType::Float, 2, sizes, buffer, dim_order, strides);
220
221 auto input_err = method->set_input(EValue(exec_aten::Tensor(&impl)), 0);
222 ASSERT_EQ(input_err, Error::Ok);
223
224 auto output_err = method->set_output_data_ptr(buffer, sizeof(buffer), 0);
225 ASSERT_EQ(output_err, Error::Ok);
226 ASSERT_EQ(method->get_output(0).toTensor().const_data_ptr(), buffer);
227
228 // Execute the method once. Cat a 1x4 to a 2x4.
229 auto execute_error = method->execute();
230 ASSERT_EQ(execute_error, Error::Ok);
231
232 auto output = method->get_output(0);
233 ASSERT_TRUE(output.isTensor());
234 EXPECT_EQ(output.toTensor().sizes()[0], 3);
235 EXPECT_EQ(output.toTensor().sizes()[1], 4);
236 // Original input should be 0.
237 for (size_t i = 0; i < 2 * 4; i++) {
238 EXPECT_FLOAT_EQ(output.toTensor().const_data_ptr<float>()[i], 0.f);
239 }
240 // Section that was cat on should be 1.
241 for (size_t i = 0; i < 1 * 4; i++) {
242 EXPECT_FLOAT_EQ(
243 output.toTensor().const_data_ptr<float>()[(2 * 4) + i], 1.f);
244 }
245
246 // Set the input again to update the size.
247 sizes[0] = output.toTensor().sizes()[0];
248 exec_aten::TensorImpl impl_2(
249 exec_aten::ScalarType::Float, 2, sizes, buffer, dim_order, strides);
250 input_err = method->set_input(EValue(exec_aten::Tensor(&impl_2)), 0);
251 ASSERT_EQ(input_err, Error::Ok);
252
253 // Execute the method again. Cat a 1x4 to a 3x4.
254 execute_error = method->execute();
255 ASSERT_EQ(execute_error, Error::Ok);
256
257 output = method->get_output(0);
258 EXPECT_EQ(output.toTensor().sizes()[0], 4);
259 EXPECT_EQ(output.toTensor().sizes()[1], 4);
260 // Original input should be 0.
261 for (size_t i = 0; i < 2 * 4; i++) {
262 EXPECT_FLOAT_EQ(output.toTensor().const_data_ptr<float>()[i], 0.f);
263 }
264 // Previous section and the new one that were cat on should be 1.
265 for (size_t i = 0; i < 2 * 4; i++) {
266 EXPECT_FLOAT_EQ(
267 output.toTensor().const_data_ptr<float>()[(2 * 4) + i], 1.f);
268 }
269 }
270
TEST_F(MethodTest,ConstantSegmentTest)271 TEST_F(MethodTest, ConstantSegmentTest) {
272 // Execute model with constants stored in segment.
273 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
274 Result<Method> method =
275 programs_["linear"]->load_method("forward", &mmm.get());
276 ASSERT_EQ(method.error(), Error::Ok);
277
278 // Can execute the method.
279 Error err = method->execute();
280 ASSERT_EQ(err, Error::Ok);
281 }
282
TEST_F(MethodTest,ConstantBufferTest)283 TEST_F(MethodTest, ConstantBufferTest) {
284 // Execute model with constants stored in the program flatbuffer.
285 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
286 Result<Method> method =
287 programs_["linear_constant_buffer"]->load_method("forward", &mmm.get());
288 ASSERT_EQ(method.error(), Error::Ok);
289
290 // Can execute the method.
291 Error err = method->execute();
292 ASSERT_EQ(err, Error::Ok);
293 }
294
295 /*
296 * TODO(T161163608): Test is disabled due to a resize bug in tensor_index_out of
297 * the portable op lib
298
299 TEST_F(MethodTest, OptionalTensorListDeserialization) {
300 ManagedMemoryManager mmm(kDefaultNonConstMemBytes,
301 kDefaultRuntimeMemBytes); Result<Method> method =
302 index_program_->load_method("forward", &mmm.get());
303 ASSERT_EQ(method.error(), Error::Ok);
304
305 // Can execute the method.
306 auto input_cleanup = prepare_input_tensors(*method);
307 ASSERT_EQ(input_cleanup.error(), Error::Ok);
308 Error err = method->execute();
309 ASSERT_EQ(err, Error::Ok);
310
311 EXPECT_EQ(method->inputs_size(), 1);
312
313 auto outputs = method->get_output(0);
314 EXPECT_EQ(outputs.toTensor().dim(), 3);
315 EXPECT_EQ(outputs.toTensor().size(0), 5);
316 EXPECT_EQ(outputs.toTensor().size(1), 2);
317 EXPECT_EQ(outputs.toTensor().size(2), 10);
318 }
319 */
320