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 #include <memory>
12
13 #include <executorch/extension/data_loader/file_data_loader.h>
14 #include <executorch/extension/runner_util/inputs.h>
15 #include <executorch/runtime/core/exec_aten/exec_aten.h>
16 #include <executorch/runtime/executor/method.h>
17 #include <executorch/runtime/executor/program.h>
18 #include <executorch/runtime/executor/test/managed_memory_manager.h>
19 #include <executorch/runtime/platform/runtime.h>
20
21 #include <gtest/gtest.h>
22
23 using namespace ::testing;
24 using exec_aten::ArrayRef;
25 using exec_aten::Scalar;
26 using exec_aten::Tensor;
27 using executorch::extension::FileDataLoader;
28 using executorch::extension::prepare_input_tensors;
29 using executorch::runtime::Error;
30 using executorch::runtime::MemoryAllocator;
31 using executorch::runtime::MemoryManager;
32 using executorch::runtime::Method;
33 using executorch::runtime::Program;
34 using executorch::runtime::Result;
35 using executorch::runtime::testing::ManagedMemoryManager;
36
37 constexpr size_t kDefaultNonConstMemBytes = 32 * 1024U;
38 constexpr size_t kDefaultRuntimeMemBytes = 32 * 1024U;
39
40 class AllocationFailureStressTest : public ::testing::Test {
41 protected:
SetUp()42 void SetUp() override {
43 executorch::runtime::runtime_init();
44
45 // Create a loader for the serialized ModuleAdd program.
46 const char* path = std::getenv("ET_MODULE_ADD_PATH");
47 Result<FileDataLoader> loader = FileDataLoader::from(path);
48 ASSERT_EQ(loader.error(), Error::Ok);
49 loader_ = std::make_unique<FileDataLoader>(std::move(loader.get()));
50
51 // Use it to load the program.
52 Result<Program> program = Program::load(
53 loader_.get(), Program::Verification::InternalConsistency);
54 ASSERT_EQ(program.error(), Error::Ok);
55 program_ = std::make_unique<Program>(std::move(program.get()));
56 }
57
58 private:
59 // Must outlive program_, but tests shouldn't need to touch it.
60 std::unique_ptr<FileDataLoader> loader_;
61
62 protected:
63 std::unique_ptr<Program> program_;
64 };
65
66 /**
67 * Slowly increases the amount of available runtime memory until load_method()
68 * and execute() succeed. This should cause every runtime allocation to fail at
69 * some point, exercising every allocation failure path reachable by the test
70 * model.
71 */
TEST_F(AllocationFailureStressTest,End2EndIncreaseRuntimeMemUntilSuccess)72 TEST_F(AllocationFailureStressTest, End2EndIncreaseRuntimeMemUntilSuccess) {
73 size_t runtime_mem_bytes = 0;
74 Error err = Error::Internal;
75 size_t num_load_failures = 0;
76 while (runtime_mem_bytes < kDefaultRuntimeMemBytes && err != Error::Ok) {
77 ManagedMemoryManager mmm(kDefaultNonConstMemBytes, runtime_mem_bytes);
78
79 // Loading should fail several times from allocation failures.
80 Result<Method> method = program_->load_method("forward", &mmm.get());
81 if (method.error() != Error::Ok) {
82 runtime_mem_bytes += sizeof(size_t);
83 num_load_failures++;
84 continue;
85 }
86
87 // Execution does not use the runtime allocator, so it should always succeed
88 // once load was successful.
89 auto input_cleanup = prepare_input_tensors(*method);
90 ASSERT_EQ(input_cleanup.error(), Error::Ok);
91 err = method->execute();
92 ASSERT_EQ(err, Error::Ok);
93 }
94 EXPECT_GT(num_load_failures, 0) << "Expected at least some failures";
95 EXPECT_EQ(err, Error::Ok)
96 << "Did not succeed after increasing runtime_mem_bytes to "
97 << runtime_mem_bytes;
98 }
99
100 /**
101 * Slowly increases the amount of available non-constant memory until
102 * load_method() and execute() succeed. This should cause every non-const
103 * allocation to fail at some point, exercising every allocation failure path
104 * reachable by the test model.
105 */
TEST_F(AllocationFailureStressTest,End2EndNonConstantMemUntilSuccess)106 TEST_F(AllocationFailureStressTest, End2EndNonConstantMemUntilSuccess) {
107 size_t non_constant_mem_bytes = 0;
108 Error err = Error::Internal;
109 size_t num_load_failures = 0;
110 while (non_constant_mem_bytes < kDefaultNonConstMemBytes &&
111 err != Error::Ok) {
112 ManagedMemoryManager mmm(non_constant_mem_bytes, kDefaultRuntimeMemBytes);
113
114 // Loading should fail several times from allocation failures.
115 Result<Method> method = program_->load_method("forward", &mmm.get());
116 if (method.error() != Error::Ok) {
117 non_constant_mem_bytes += sizeof(size_t);
118 num_load_failures++;
119 continue;
120 }
121
122 // Execution does not use the runtime allocator, so it should always succeed
123 // once load was successful.
124 auto input_cleanup = prepare_input_tensors(*method);
125 ASSERT_EQ(input_cleanup.error(), Error::Ok);
126 err = method->execute();
127 ASSERT_EQ(err, Error::Ok);
128 }
129 EXPECT_GT(num_load_failures, 0) << "Expected at least some failures";
130 EXPECT_EQ(err, Error::Ok)
131 << "Did not succeed after increasing non_constant_mem_bytes to "
132 << non_constant_mem_bytes;
133 }
134