xref: /aosp_15_r20/external/executorch/extension/memory_allocator/malloc_memory_allocator.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
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