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 <stdio.h> 12 #include <cinttypes> 13 #include <cstdint> 14 15 #include <executorch/runtime/core/error.h> 16 #include <executorch/runtime/platform/assert.h> 17 #include <executorch/runtime/platform/compiler.h> 18 #include <executorch/runtime/platform/log.h> 19 #include <executorch/runtime/platform/profiler.h> 20 21 namespace executorch { 22 namespace runtime { 23 24 /** 25 * A class that does simple allocation based on a size and returns the pointer 26 * to the memory address. It bookmarks a buffer with certain size. The 27 * allocation is simply checking space and growing the cur_ pointer with each 28 * allocation request. 29 * 30 * Simple example: 31 * 32 * // User allocates a 100 byte long memory in the heap. 33 * uint8_t* memory_pool = malloc(100 * sizeof(uint8_t)); 34 * MemoryAllocator allocator(100, memory_pool) 35 * // Pass allocator object in the Executor 36 * 37 * Underneath the hood, ExecuTorch will call 38 * allocator.allocate() to keep iterating cur_ pointer 39 */ 40 class MemoryAllocator { 41 public: 42 /** 43 * Default alignment of memory returned by this class. Ensures that pointer 44 * fields of structs will be aligned. Larger types like `long double` may not 45 * be, however, depending on the toolchain and architecture. 46 */ 47 static constexpr size_t kDefaultAlignment = alignof(void*); 48 49 /** 50 * Constructs a new memory allocator of a given `size`, starting at the 51 * provided `base_address`. 52 * 53 * @param[in] size The size in bytes of the buffer at `base_address`. 54 * @param[in] base_address The buffer to allocate from. Does not take 55 * ownership of this buffer, so it must be valid for the lifetime of of 56 * the MemoryAllocator. 57 */ MemoryAllocator(uint32_t size,uint8_t * base_address)58 MemoryAllocator(uint32_t size, uint8_t* base_address) 59 : begin_(base_address), 60 end_(base_address + size), 61 cur_(base_address), 62 size_(size) {} 63 64 /** 65 * Allocates `size` bytes of memory. 66 * 67 * @param[in] size Number of bytes to allocate. 68 * @param[in] alignment Minimum alignment for the returned pointer. Must be a 69 * power of 2. 70 * 71 * @returns Aligned pointer to the allocated memory on success. 72 * @retval nullptr Not enough memory, or `alignment` was not a power of 2. 73 */ 74 virtual void* allocate(size_t size, size_t alignment = kDefaultAlignment) { 75 if (!isPowerOf2(alignment)) { 76 ET_LOG(Error, "Alignment %zu is not a power of 2", alignment); 77 return nullptr; 78 } 79 80 // The allocation will occupy [start, end), where the start is the next 81 // position that's a multiple of alignment. 82 uint8_t* start = alignPointer(cur_, alignment); 83 uint8_t* end = start + size; 84 85 // If the end of this allocation exceeds the end of this allocator, print 86 // error messages and return nullptr 87 if (end > end_) { 88 ET_LOG( 89 Error, 90 "Memory allocation failed: %zuB requested (adjusted for alignment), %zuB available", 91 static_cast<size_t>(end - cur_), 92 static_cast<size_t>(end_ - cur_)); 93 return nullptr; 94 } 95 96 // Otherwise, record how many bytes were used, advance cur_ to the new end, 97 // and then return start. Note that the number of bytes used is (end - cur_) 98 // instead of (end - start) because start > cur_ if there is a misalignment 99 EXECUTORCH_TRACK_ALLOCATION(prof_id_, end - cur_); 100 cur_ = end; 101 return static_cast<void*>(start); 102 } 103 104 /** 105 * Allocates a buffer large enough for an instance of type T. Note that the 106 * memory will not be initialized. 107 * 108 * Example: 109 * @code 110 * auto p = memory_allocator->allocateInstance<MyType>(); 111 * @endcode 112 * 113 * @param[in] alignment Minimum alignment for the returned pointer. Must be a 114 * power of 2. Defaults to the natural alignment of T. 115 * 116 * @returns Aligned pointer to the allocated memory on success. 117 * @retval nullptr Not enough memory, or `alignment` was not a power of 2. 118 */ 119 template <typename T> 120 T* allocateInstance(size_t alignment = alignof(T)) { 121 return static_cast<T*>(this->allocate(sizeof(T), alignment)); 122 } 123 124 /** 125 * Allocates `size` number of chunks of type T, where each chunk is of size 126 * equal to sizeof(T) bytes. 127 * 128 * @param[in] size Number of memory chunks to allocate. 129 * @param[in] alignment Minimum alignment for the returned pointer. Must be a 130 * power of 2. Defaults to the natural alignment of T. 131 * 132 * @returns Aligned pointer to the allocated memory on success. 133 * @retval nullptr Not enough memory, or `alignment` was not a power of 2. 134 */ 135 template <typename T> 136 T* allocateList(size_t size, size_t alignment = alignof(T)) { 137 // Some users of this method allocate lists of pointers, causing the next 138 // line to expand to `sizeof(type *)`, which triggers a clang-tidy warning. 139 // NOLINTNEXTLINE(bugprone-sizeof-expression) 140 return static_cast<T*>(this->allocate(size * sizeof(T), alignment)); 141 } 142 143 // Returns the allocator memory's base address. base_address()144 virtual uint8_t* base_address() const { 145 return begin_; 146 } 147 148 // Returns the total size of the allocator's memory buffer. size()149 virtual uint32_t size() const { 150 return size_; 151 } 152 153 // Resets the current pointer to the base address. It does nothing to 154 // the contents. reset()155 virtual void reset() { 156 cur_ = begin_; 157 } 158 enable_profiling(ET_UNUSED const char * name)159 void enable_profiling(ET_UNUSED const char* name) { 160 prof_id_ = EXECUTORCH_TRACK_ALLOCATOR(name); 161 } 162 ~MemoryAllocator()163 virtual ~MemoryAllocator() {} 164 165 protected: 166 /** 167 * Returns the profiler ID for this allocator. 168 */ prof_id()169 int32_t prof_id() const { 170 return prof_id_; 171 } 172 173 /** 174 * Returns true if the value is an integer power of 2. 175 */ isPowerOf2(size_t value)176 static bool isPowerOf2(size_t value) { 177 return value > 0 && (value & ~(value - 1)) == value; 178 } 179 180 /** 181 * Returns the next alignment for a given pointer. 182 */ alignPointer(void * ptr,size_t alignment)183 static uint8_t* alignPointer(void* ptr, size_t alignment) { 184 intptr_t addr = reinterpret_cast<intptr_t>(ptr); 185 if ((addr & (alignment - 1)) == 0) { 186 // Already aligned. 187 return reinterpret_cast<uint8_t*>(ptr); 188 } 189 addr = (addr | (alignment - 1)) + 1; 190 return reinterpret_cast<uint8_t*>(addr); 191 } 192 193 private: 194 uint8_t* const begin_; 195 uint8_t* const end_; 196 uint8_t* cur_; 197 uint32_t const size_; 198 int32_t prof_id_ = -1; 199 }; 200 201 #if ET_HAVE_GNU_STATEMENT_EXPRESSIONS 202 /** 203 * Tries allocating from the specified MemoryAllocator*. 204 * 205 * - On success, returns a pointer to the allocated buffer. 206 * - On failure, executes the provided code block, which must return or panic. 207 * 208 * Example: 209 * @code 210 * char* buf = ET_TRY_ALLOCATE_OR( 211 * memory_allocator, bufsize, { 212 * *out_err = Error::MemoryAllocationFailed; 213 * return nullopt; 214 * }); 215 * @endcode 216 */ 217 #define ET_TRY_ALLOCATE_OR(memory_allocator__, nbytes__, ...) \ 218 ({ \ 219 void* et_try_allocate_result = memory_allocator__->allocate(nbytes__); \ 220 if (et_try_allocate_result == nullptr && nbytes__ > 0) { \ 221 __VA_ARGS__ \ 222 /* The args must return. */ \ 223 ET_UNREACHABLE(); \ 224 } \ 225 et_try_allocate_result; \ 226 }) 227 228 /** 229 * Tries allocating an instance of type__ from the specified MemoryAllocator*. 230 * 231 * - On success, returns a pointer to the allocated buffer. Note that the memory 232 * will not be initialized. 233 * - On failure, executes the provided code block, which must return or panic. 234 * 235 * Example: 236 * @code 237 * char* buf = ET_TRY_ALLOCATE_INSTANCE_OR( 238 * memory_allocator, 239 * MyType, 240 * { *out_err = Error::MemoryAllocationFailed; return nullopt; }); 241 * @endcode 242 */ 243 #define ET_TRY_ALLOCATE_INSTANCE_OR(memory_allocator__, type__, ...) \ 244 ({ \ 245 type__* et_try_allocate_result = \ 246 memory_allocator__->allocateInstance<type__>(); \ 247 if (et_try_allocate_result == nullptr) { \ 248 __VA_ARGS__ \ 249 /* The args must return. */ \ 250 ET_UNREACHABLE(); \ 251 } \ 252 et_try_allocate_result; \ 253 }) 254 255 /** 256 * Tries allocating multiple elements of a given type from the specified 257 * MemoryAllocator*. 258 * 259 * - On success, returns a pointer to the allocated buffer. 260 * - On failure, executes the provided code block, which must return or panic. 261 * 262 * Example: 263 * @code 264 * Tensor* tensor_list = ET_TRY_ALLOCATE_LIST_OR( 265 * memory_allocator, Tensor, num_tensors, { 266 * *out_err = Error::MemoryAllocationFailed; 267 * return nullopt; 268 * }); 269 * @endcode 270 */ 271 #define ET_TRY_ALLOCATE_LIST_OR(memory_allocator__, type__, nelem__, ...) \ 272 ({ \ 273 type__* et_try_allocate_result = \ 274 memory_allocator__->allocateList<type__>(nelem__); \ 275 if (et_try_allocate_result == nullptr && nelem__ > 0) { \ 276 __VA_ARGS__ \ 277 /* The args must return. */ \ 278 ET_UNREACHABLE(); \ 279 } \ 280 et_try_allocate_result; \ 281 }) 282 #else // !ET_HAVE_GNU_STATEMENT_EXPRESSIONS 283 /** 284 * The recommended alternative for statement expression-incompatible compilers 285 * is to directly allocate the memory. 286 * e.g. memory_allocator__->allocate(nbytes__); 287 */ 288 #define ET_TRY_ALLOCATE_OR(memory_allocator__, nbytes__, ...) \ 289 static_assert( \ 290 false, \ 291 "ET_TRY_ALLOCATE_OR uses statement expressions and \ 292 thus is not available for use with this compiler."); 293 294 /** 295 * The recommended alternative for statement expression-incompatible compilers 296 * is to directly allocate the memory. 297 * e.g. memory_allocator__->allocateInstance<type__>(); 298 */ 299 #define ET_TRY_ALLOCATE_INSTANCE_OR(memory_allocator__, type__, ...) \ 300 static_assert( \ 301 false, \ 302 "ET_TRY_ALLOCATE_INSTANCE_OR uses statement \ 303 expressions and thus is not available for use with this compiler."); 304 305 /** 306 * The recommended alternative for statement expression-incompatible compilers 307 * is to directly use allocate the memory. 308 * e.g. memory_allocator__->allocateList<type__>(nelem__); 309 */ 310 #define ET_TRY_ALLOCATE_LIST_OR(memory_allocator__, type__, nelem__, ...) \ 311 static_assert( \ 312 false, \ 313 "ET_TRY_ALLOCATE_LIST_OR uses statement \ 314 expressions and thus is not available for use with this compiler."); 315 #endif // !ET_HAVE_GNU_STATEMENT_EXPRESSIONS 316 317 /** 318 * Tries allocating from the specified MemoryAllocator*. 319 * 320 * - On success, returns a pointer to the allocated buffer. 321 * - On failure, returns `Error::MemoryAllocationFailed` from the calling 322 * function, which must be declared to return `executorch::runtime::Error`. 323 * 324 * Example: 325 * @code 326 * char* buf = ET_ALLOCATE_OR_RETURN_ERROR(memory_allocator, bufsize); 327 * @endcode 328 */ 329 #define ET_ALLOCATE_OR_RETURN_ERROR(memory_allocator__, nbytes__) \ 330 ET_TRY_ALLOCATE_OR(memory_allocator__, nbytes__, { \ 331 return ::executorch::runtime::Error::MemoryAllocationFailed; \ 332 }) 333 334 /** 335 * Tries allocating an instance of type__ from the specified MemoryAllocator*. 336 * 337 * - On success, returns a pointer to the allocated buffer. Note that the memory 338 * will not be initialized. 339 * - On failure, returns `Error::MemoryAllocationFailed` from the calling 340 * function, which must be declared to return `executorch::runtime::Error`. 341 * 342 * Example: 343 * @code 344 * char* buf = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(memory_allocator, MyType); 345 * @endcode 346 */ 347 #define ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(memory_allocator__, type__) \ 348 ET_TRY_ALLOCATE_INSTANCE_OR(memory_allocator__, type__, { \ 349 return ::executorch::runtime::Error::MemoryAllocationFailed; \ 350 }) 351 352 /** 353 * Tries allocating multiple elements of a given type from the specified 354 * MemoryAllocator*. 355 * 356 * - On success, returns a pointer to the allocated buffer. 357 * - On failure, returns `Error::MemoryAllocationFailed` from the calling 358 * function, which must be declared to return `executorch::runtime::Error`. 359 * 360 * Example: 361 * @code 362 * Tensor* tensor_list = ET_ALLOCATE_LIST_OR_RETURN_ERROR( 363 * memory_allocator, Tensor, num_tensors); 364 * @endcode 365 */ 366 #define ET_ALLOCATE_LIST_OR_RETURN_ERROR(memory_allocator__, type__, nelem__) \ 367 ET_TRY_ALLOCATE_LIST_OR(memory_allocator__, type__, nelem__, { \ 368 return ::executorch::runtime::Error::MemoryAllocationFailed; \ 369 }) 370 371 } // namespace runtime 372 } // namespace executorch 373 374 namespace torch { 375 namespace executor { 376 // TODO(T197294990): Remove these deprecated aliases once all users have moved 377 // to the new `::executorch` namespaces. 378 using ::executorch::runtime::MemoryAllocator; 379 } // namespace executor 380 } // namespace torch 381