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 <cstddef> 12 #include <cstdint> 13 #include <vector> 14 15 #include <executorch/runtime/core/memory_allocator.h> 16 17 namespace executorch { 18 namespace extension { 19 20 /** 21 * Dynamically allocates memory using malloc() and frees all pointers at 22 * destruction time. 23 * 24 * For systems with malloc(), this can be easier than using a fixed-sized 25 * MemoryAllocator. 26 */ 27 class MallocMemoryAllocator : public executorch::runtime::MemoryAllocator { 28 public: 29 /** 30 * Construct a new Malloc memory allocator via an optional alignment size 31 * parameter. 32 * 33 * @param[in] align_size An optional parameter to specify alignment parameter 34 * for each allocate() call. 35 */ MallocMemoryAllocator()36 MallocMemoryAllocator() : MemoryAllocator(0, nullptr) {} 37 ~MallocMemoryAllocator()38 ~MallocMemoryAllocator() override { 39 reset(); 40 } 41 42 /** 43 * Allocates 'size' bytes of memory, returning a pointer to the allocated 44 * region, or nullptr upon failure. The size will be rounded up based on the 45 * memory alignment size. 46 */ 47 void* allocate(size_t size, size_t alignment = kDefaultAlignment) override { 48 EXECUTORCH_TRACK_ALLOCATION(prof_id(), size); 49 50 if (!isPowerOf2(alignment)) { 51 ET_LOG(Error, "Alignment %zu is not a power of 2", alignment); 52 return nullptr; 53 } 54 55 // The minimum alignment that malloc() is guaranteed to provide. 56 static constexpr size_t kMallocAlignment = alignof(std::max_align_t); 57 if (alignment > kMallocAlignment) { 58 // To get higher alignments, allocate extra and then align the returned 59 // pointer. This will waste an extra `alignment` bytes every time, but 60 // this is the only portable way to get aligned memory from the heap. 61 size += alignment; 62 } 63 mem_ptrs_.emplace_back(std::malloc(size)); 64 return alignPointer(mem_ptrs_.back(), alignment); 65 } 66 67 // Free up each hosted memory pointer. The memory was created via malloc. reset()68 void reset() override { 69 for (auto mem_ptr : mem_ptrs_) { 70 free(mem_ptr); 71 } 72 mem_ptrs_.clear(); 73 } 74 75 private: 76 std::vector<void*> mem_ptrs_; 77 }; 78 79 } // namespace extension 80 } // namespace executorch 81 82 namespace torch { 83 namespace executor { 84 namespace util { 85 // TODO(T197294990): Remove these deprecated aliases once all users have moved 86 // to the new `::executorch` namespaces. 87 using ::executorch::extension::MallocMemoryAllocator; 88 } // namespace util 89 } // namespace executor 90 } // namespace torch 91