|  | 
|  | 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 <executorch/extension/data_loader/file_data_loader.h> | 
|  | 10 | + #include <executorch/extension/flat_tensor/flat_tensor_data_map.h> | 
|  | 11 | + #include <executorch/runtime/core/error.h> | 
|  | 12 | + #include <executorch/runtime/core/result.h> | 
|  | 13 | + #include <executorch/runtime/executor/method.h> | 
|  | 14 | + #include <executorch/runtime/executor/program.h> | 
|  | 15 | + #include <executorch/runtime/executor/test/managed_memory_manager.h> | 
|  | 16 | + #include <executorch/runtime/platform/runtime.h> | 
|  | 17 | +  | 
|  | 18 | + #include <gtest/gtest.h> | 
|  | 19 | +  | 
|  | 20 | + using namespace ::testing; | 
|  | 21 | + using executorch::extension::FlatTensorDataMap; | 
|  | 22 | + using executorch::runtime::DataLoader; | 
|  | 23 | + using executorch::runtime::Error; | 
|  | 24 | + using executorch::runtime::FreeableBuffer; | 
|  | 25 | + using executorch::runtime::Method; | 
|  | 26 | + using executorch::runtime::Program; | 
|  | 27 | + using executorch::runtime::Result; | 
|  | 28 | + using executorch::runtime::testing::ManagedMemoryManager; | 
|  | 29 | + using torch::executor::util::FileDataLoader; | 
|  | 30 | +  | 
|  | 31 | + constexpr size_t kDefaultNonConstMemBytes = 32 * 1024U; | 
|  | 32 | + constexpr size_t kDefaultRuntimeMemBytes = 32 * 1024U; | 
|  | 33 | +  | 
|  | 34 | + class DataSeparationTest : public ::testing::Test { | 
|  | 35 | +  protected: | 
|  | 36 | +   void SetUp() override { | 
|  | 37 | +     // Since these tests cause ET_LOG to be called, the PAL must be initialized | 
|  | 38 | +     // first. | 
|  | 39 | +     executorch::runtime::runtime_init(); | 
|  | 40 | +  | 
|  | 41 | +     // Create data loaders. | 
|  | 42 | +     Result<FileDataLoader> linear_program_loader = | 
|  | 43 | +         FileDataLoader::from(std::getenv("ET_MODULE_LINEAR_XNN_PROGRAM_PATH")); | 
|  | 44 | +     ASSERT_EQ(linear_program_loader.error(), Error::Ok); | 
|  | 45 | +     linear_program_loader_ = std::make_unique<FileDataLoader>( | 
|  | 46 | +         std::move(linear_program_loader.get())); | 
|  | 47 | +  | 
|  | 48 | +     Result<FileDataLoader> linear_data_loader = | 
|  | 49 | +         FileDataLoader::from(std::getenv("ET_MODULE_LINEAR_XNN_DATA_PATH")); | 
|  | 50 | +     ASSERT_EQ(linear_data_loader.error(), Error::Ok); | 
|  | 51 | +     linear_data_loader_ = | 
|  | 52 | +         std::make_unique<FileDataLoader>(std::move(linear_data_loader.get())); | 
|  | 53 | +  | 
|  | 54 | +     // Create programs. | 
|  | 55 | +     Result<Program> linear_program = Program::load( | 
|  | 56 | +         linear_program_loader_.get(), | 
|  | 57 | +         Program::Verification::InternalConsistency); | 
|  | 58 | +     ASSERT_EQ(linear_program.error(), Error::Ok); | 
|  | 59 | +     linear_program_ = | 
|  | 60 | +         std::make_unique<Program>(std::move(linear_program.get())); | 
|  | 61 | +  | 
|  | 62 | +     Result<FlatTensorDataMap> linear_data_map = | 
|  | 63 | +         FlatTensorDataMap::load(linear_data_loader_.get()); | 
|  | 64 | +     EXPECT_EQ(linear_data_map.error(), Error::Ok); | 
|  | 65 | +     linear_data_map_ = | 
|  | 66 | +         std::make_unique<FlatTensorDataMap>(std::move(linear_data_map.get())); | 
|  | 67 | +   } | 
|  | 68 | +  | 
|  | 69 | +  private: | 
|  | 70 | +   std::unique_ptr<FileDataLoader> linear_program_loader_; | 
|  | 71 | +   std::unique_ptr<FileDataLoader> linear_data_loader_; | 
|  | 72 | +  | 
|  | 73 | +  protected: | 
|  | 74 | +   std::unique_ptr<Program> linear_program_; | 
|  | 75 | +   std::unique_ptr<FlatTensorDataMap> linear_data_map_; | 
|  | 76 | + }; | 
|  | 77 | +  | 
|  | 78 | + TEST_F(DataSeparationTest, TestExternalData) { | 
|  | 79 | +    FlatTensorDataMap* data_map = linear_data_map_.get(); | 
|  | 80 | +    EXPECT_EQ(data_map->get_num_keys().get(), 2); | 
|  | 81 | + | 
|  | 82 | +    Result<const char*> key0 = data_map->get_key(0); | 
|  | 83 | +    EXPECT_EQ(key0.error(), Error::Ok); | 
|  | 84 | +    Result<const char*> key1 = data_map->get_key(1); | 
|  | 85 | +    EXPECT_EQ(key1.error(), Error::Ok); | 
|  | 86 | + | 
|  | 87 | +    // Check that accessing keys out of bounds fails. | 
|  | 88 | +    EXPECT_EQ(data_map->get_key(2).error(), Error::InvalidArgument); | 
|  | 89 | + | 
|  | 90 | +    // Linear.weight | 
|  | 91 | +    Result<FreeableBuffer> data0 = data_map->get_data(key0.get()); | 
|  | 92 | +    EXPECT_EQ(data0.error(), Error::Ok); | 
|  | 93 | +    EXPECT_EQ(data0.get().size(), 36); // 3*3*4 (3*3 matrix, 4 bytes per float) | 
|  | 94 | + | 
|  | 95 | +    // Linear.bias | 
|  | 96 | +    Result<FreeableBuffer> data1 = data_map->get_data(key1.get()); | 
|  | 97 | +    EXPECT_EQ(data1.error(), Error::Ok); | 
|  | 98 | +    EXPECT_EQ(data1.get().size(), 12); // 3*4 (3 vector, 4 bytes per float) | 
|  | 99 | + | 
|  | 100 | +    // Check that accessing non-existent data fails. | 
|  | 101 | +    Result<FreeableBuffer> data2 = data_map->get_data("nonexistent"); | 
|  | 102 | +    EXPECT_EQ(data2.error(), Error::NotFound); | 
|  | 103 | +} | 
|  | 104 | + | 
|  | 105 | + TEST_F(DataSeparationTest, TestE2E) { | 
|  | 106 | +   ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); | 
|  | 107 | +   Result<Method> method = linear_program_->load_method( | 
|  | 108 | +       "forward", &mmm.get(), nullptr, linear_data_map_.get()); | 
|  | 109 | +   ASSERT_EQ(method.error(), Error::Ok); | 
|  | 110 | +  | 
|  | 111 | +   // Can execute the method. | 
|  | 112 | +   Error err = method->execute(); | 
|  | 113 | +   ASSERT_EQ(err, Error::Ok); | 
|  | 114 | + } | 
0 commit comments