xref: /aosp_15_r20/external/armnn/src/backends/backendsCommon/MemoryManager.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 #pragma once
6 
7 #include <armnn/backends/ICustomAllocator.hpp>
8 
9 namespace armnn
10 {
11 struct Allocator
12 {
13     /// Pointer to @ICustomAllocator.
14     std::shared_ptr<ICustomAllocator> m_CustomAllocator{};
15     /// Value which the size of each buffer (actual data size + padding) has to be a multiple of.
16     size_t m_Alignment = 0 ;
17 };
18 
19 struct TensorMemory
20 {
21     /// Number of bytes the value is away from the @BufferStorage.m_Buffer.
22     size_t m_Offset{};
23     /// Identifier to be used by the @LoadedNetwork to order the tensors.
24     unsigned int m_OutputSlotId{};
25     /// Pointer to the tensor value.
26     void* m_Data = nullptr;
27 };
28 
29 struct BufferStorage
30 {
31     /// Vector of pointer to @TensorMemory.
32     std::vector<std::shared_ptr<TensorMemory>> m_TensorMemoryVector;
33     /// Total size of the buffer.
34     size_t m_BufferSize;
35     /// Pointer to the first element of the buffer.
36     void* m_Buffer = nullptr;
37 };
38 
39 class MemoryManager
40 {
41 public:
42     /// Initialization method to store in m_AllocatorBufferStoragePairVector all information needed.
43     /// @param[in] bufferStorageVector - Vector of BufferStorage.
44     /// @param[in] customAllocator - Pointer to ICustomAllocator.
45     /// @param[in] typeAlignment - Optional parameter. Value of which the size of each value has to be multiple of.
46     void StoreMemToAllocate(std::vector<BufferStorage> bufferStorageVector,
47                             std::shared_ptr<ICustomAllocator> customAllocator,
48                             size_t typeAlignment = 0);
49 
50     /// Allocate the amount of memory indicated by m_BufferSize, and
51     /// point each m_Data to each correspondent Tensor so that they are m_Offset bytes separated.
52     void Allocate();
53 
54     /// Deallocate memory
55     void Deallocate();
56 
57 private:
58     std::vector<std::pair<Allocator, std::vector<BufferStorage>>> m_AllocatorBufferStoragePairVector;
59 };
60 
61 } // namespace armnn
62