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 #pragma once 10 11 #include <executorch/runtime/core/data_loader.h> 12 #include <executorch/runtime/core/result.h> 13 #include <executorch/runtime/platform/compiler.h> 14 15 namespace executorch { 16 namespace extension { 17 18 /** 19 * A DataLoader that loads segments from a file, allocating the memory 20 * with `malloc()`. 21 * 22 * Note that this will keep the file open for the duration of its lifetime, to 23 * avoid the overhead of opening it again for every load() call. 24 */ 25 class MmapDataLoader final : public executorch::runtime::DataLoader { 26 public: 27 /** 28 * Describes how and whether to lock loaded pages with `mlock()`. 29 * 30 * Using `mlock()` typically loads all of the pages immediately, and will 31 * typically ensure that they are not swapped out. The actual behavior 32 * will depend on the host system. 33 */ 34 enum class MlockConfig { 35 /// Do not call `mlock()` on loaded pages. 36 NoMlock, 37 /// Call `mlock()` on loaded pages, failing if it fails. 38 UseMlock, 39 /// Call `mlock()` on loaded pages, ignoring errors if it fails. 40 UseMlockIgnoreErrors, 41 }; 42 43 /** 44 * Creates a new MmapDataLoader that wraps the named file. Fails if 45 * the file can't be opened for reading or if its size can't be found. 46 * 47 * @param[in] file_name The path to the file to load from. The file will be 48 * kept open until the MmapDataLoader is destroyed, to avoid the 49 * overhead of opening it again for every load() call. 50 * @param[in] mlock_config How and whether to lock loaded pages with 51 * `mlock()`. 52 */ 53 static executorch::runtime::Result<MmapDataLoader> from( 54 const char* file_name, 55 MlockConfig mlock_config = MlockConfig::UseMlock); 56 57 /// DEPRECATED: Use the lowercase `from()` instead. 58 ET_DEPRECATED static executorch::runtime::Result<MmapDataLoader> From( 59 const char* file_name, 60 MlockConfig mlock_config = MlockConfig::UseMlock) { 61 return from(file_name, mlock_config); 62 } 63 64 /// DEPRECATED: Use the version of `from()` that takes an MlockConfig. 65 ET_DEPRECATED From(const char * file_name,bool use_mlock)66 static executorch::runtime::Result<MmapDataLoader> From( 67 const char* file_name, 68 bool use_mlock) { 69 MlockConfig mlock_config = 70 use_mlock ? MlockConfig::UseMlock : MlockConfig::NoMlock; 71 return from(file_name, mlock_config); 72 } 73 74 // Movable to be compatible with Result. MmapDataLoader(MmapDataLoader && rhs)75 MmapDataLoader(MmapDataLoader&& rhs) noexcept 76 : file_name_(rhs.file_name_), 77 file_size_(rhs.file_size_), 78 page_size_(rhs.page_size_), 79 fd_(rhs.fd_), 80 mlock_config_(rhs.mlock_config_) { 81 const_cast<const char*&>(rhs.file_name_) = nullptr; 82 const_cast<size_t&>(rhs.file_size_) = 0; 83 const_cast<size_t&>(rhs.page_size_) = 0; 84 const_cast<int&>(rhs.fd_) = -1; 85 const_cast<MlockConfig&>(rhs.mlock_config_) = MlockConfig::NoMlock; 86 } 87 88 ~MmapDataLoader() override; 89 90 ET_NODISCARD 91 executorch::runtime::Result<executorch::runtime::FreeableBuffer> load( 92 size_t offset, 93 size_t size, 94 const DataLoader::SegmentInfo& segment_info) const override; 95 96 ET_NODISCARD executorch::runtime::Result<size_t> size() const override; 97 98 private: MmapDataLoader(int fd,size_t file_size,const char * file_name,size_t page_size,MlockConfig mlock_config)99 MmapDataLoader( 100 int fd, 101 size_t file_size, 102 const char* file_name, 103 size_t page_size, 104 MlockConfig mlock_config) 105 : file_name_(file_name), 106 file_size_(file_size), 107 page_size_(page_size), 108 fd_(fd), 109 mlock_config_(mlock_config) {} 110 111 // Not safely copyable. 112 MmapDataLoader(const MmapDataLoader&) = delete; 113 MmapDataLoader& operator=(const MmapDataLoader&) = delete; 114 MmapDataLoader& operator=(MmapDataLoader&&) = delete; 115 116 const char* const file_name_; // String data is owned by the instance. 117 const size_t file_size_; 118 const size_t page_size_; 119 const int fd_; // Owned by the instance. 120 const MlockConfig mlock_config_; 121 }; 122 123 } // namespace extension 124 } // namespace executorch 125 126 namespace torch { 127 namespace executor { 128 namespace util { 129 // TODO(T197294990): Remove these deprecated aliases once all users have moved 130 // to the new `::executorch` namespaces. 131 using ::executorch::extension::MmapDataLoader; 132 } // namespace util 133 } // namespace executor 134 } // namespace torch 135