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/memory_allocator.h> 12 #include <executorch/runtime/core/result.h> 13 #include <executorch/runtime/core/span.h> 14 #include <executorch/runtime/platform/assert.h> 15 #include <executorch/runtime/platform/compiler.h> 16 #include <executorch/runtime/platform/log.h> 17 #include <cstdint> 18 19 namespace executorch { 20 namespace runtime { 21 22 /** 23 * A group of buffers that can be used to represent a device's memory hierarchy. 24 */ 25 class HierarchicalAllocator final { 26 public: 27 /** 28 * Constructs a new hierarchical allocator with the given array of buffers. 29 * 30 * - Memory IDs are based on the index into `buffers`: `buffers[N]` will have 31 * a memory ID of `N`. 32 * - `buffers.size()` must be >= `MethodMeta::num_non_const_buffers()`. 33 * - `buffers[N].size()` must be >= `MethodMeta::non_const_buffer_size(N)`. 34 */ HierarchicalAllocator(Span<Span<uint8_t>> buffers)35 explicit HierarchicalAllocator(Span<Span<uint8_t>> buffers) 36 : buffers_(buffers) {} 37 38 /** 39 * DEPRECATED: Use spans instead. 40 */ HierarchicalAllocator(uint32_t n_allocators,MemoryAllocator * allocators)41 ET_DEPRECATED HierarchicalAllocator( 42 uint32_t n_allocators, 43 MemoryAllocator* allocators) 44 : buffers_(to_spans(n_allocators, allocators)) {} 45 46 /** 47 * Returns the address at the byte offset `offset_bytes` from the given 48 * buffer's base address, which points to at least `size_bytes` of memory. 49 * 50 * @param[in] memory_id The ID of the buffer in the hierarchy. 51 * @param[in] offset_bytes The offset in bytes into the specified buffer. 52 * @param[in] size_bytes The amount of memory that should be available at 53 * the offset. 54 * 55 * @returns On success, the address of the requested byte offset into the 56 * specified buffer. On failure, a non-Ok Error. 57 */ get_offset_address(uint32_t memory_id,size_t offset_bytes,size_t size_bytes)58 ET_NODISCARD Result<void*> get_offset_address( 59 uint32_t memory_id, 60 size_t offset_bytes, 61 size_t size_bytes) { 62 ET_CHECK_OR_RETURN_ERROR( 63 memory_id < buffers_.size(), 64 InvalidArgument, 65 "id %" PRIu32 " >= %zu", 66 memory_id, 67 buffers_.size()); 68 Span<uint8_t> buffer = buffers_[memory_id]; 69 ET_CHECK_OR_RETURN_ERROR( 70 offset_bytes + size_bytes <= buffer.size(), 71 MemoryAllocationFailed, 72 "offset_bytes (%zu) + size_bytes (%zu) >= allocator size (%zu) " 73 "for memory_id %" PRIu32, 74 offset_bytes, 75 size_bytes, 76 buffer.size(), 77 memory_id); 78 return buffer.data() + offset_bytes; 79 } 80 81 private: 82 // TODO(T162089316): Remove the span array and to_spans once all users move to 83 // spans. This array is necessary to hold the pointers and sizes that were 84 // originally provided as MemoryAllocator instances. 85 static constexpr size_t kSpanArraySize = 16; 86 // NOTE: span_array_ must be declared before buffers_ so that it isn't 87 // re-initialized to zeros after initializing buffers_. 88 Span<uint8_t> span_array_[kSpanArraySize]; to_spans(uint32_t n_allocators,MemoryAllocator * allocators)89 Span<Span<uint8_t>> to_spans( 90 uint32_t n_allocators, 91 MemoryAllocator* allocators) { 92 ET_CHECK_MSG( 93 n_allocators <= kSpanArraySize, 94 "n_allocators %" PRIu32 " > %zu", 95 n_allocators, 96 kSpanArraySize); 97 for (uint32_t i = 0; i < n_allocators; ++i) { 98 span_array_[i] = 99 Span<uint8_t>(allocators[i].base_address(), allocators[i].size()); 100 } 101 return {span_array_, n_allocators}; 102 } 103 104 /// The underlying buffers. 105 Span<Span<uint8_t>> buffers_; 106 }; 107 108 } // namespace runtime 109 } // namespace executorch 110 111 namespace torch { 112 namespace executor { 113 // TODO(T197294990): Remove these deprecated aliases once all users have moved 114 // to the new `::executorch` namespaces. 115 using ::executorch::runtime::HierarchicalAllocator; 116 } // namespace executor 117 } // namespace torch 118