1 /*
2  * Copyright (C) 2020 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef ANDROID_PACKAGES_MODULES_NEURALNETWORKS_COMMON_TYPES_NNAPI_SHARED_MEMORY_H
18 #define ANDROID_PACKAGES_MODULES_NEURALNETWORKS_COMMON_TYPES_NNAPI_SHARED_MEMORY_H
19 
20 #include <android-base/unique_fd.h>
21 
22 #include <any>
23 #include <memory>
24 #include <optional>
25 #include <string>
26 #include <utility>
27 #include <variant>
28 #include <vector>
29 
30 #include "nnapi/Result.h"
31 #include "nnapi/Types.h"
32 
33 namespace android::nn {
34 
35 class MutableMemoryBuilder {
36    public:
37     explicit MutableMemoryBuilder(uint32_t poolIndex);
38 
39     DataLocation append(size_t length, size_t alignment = kMinMemoryAlignment,
40                         size_t padding = kMinMemoryPadding);
41     bool empty() const;
42 
43     GeneralResult<SharedMemory> finish();
44 
45    private:
46     uint32_t mPoolIndex;
47     size_t mSize = 0;
48 };
49 
50 class ConstantMemoryBuilder {
51    public:
52     explicit ConstantMemoryBuilder(uint32_t poolIndex);
53 
54     DataLocation append(const void* data, size_t length);
55     bool empty() const;
56 
57     GeneralResult<SharedMemory> finish();
58 
59    private:
60     struct LazyCopy {
61         const void* data;
62         size_t length;
63         size_t offset;
64     };
65 
66     MutableMemoryBuilder mBuilder;
67     std::vector<LazyCopy> mSlices;
68 };
69 
70 GeneralResult<base::unique_fd> dupFd(int fd);
71 
72 // Precondition: `*ForwardFdIt` must be convertible to `int`
73 template <typename ForwardFdIt>
dupFds(ForwardFdIt first,ForwardFdIt last)74 GeneralResult<std::vector<base::unique_fd>> dupFds(ForwardFdIt first, ForwardFdIt last) {
75     std::vector<base::unique_fd> fds;
76     fds.reserve(std::distance(first, last));
77     for (; first != last; ++first) {
78         const int fd = *first;
79         fds.push_back(NN_TRY(dupFd(fd)));
80     }
81     return fds;
82 }
83 
84 // Precondition: size > 0
85 GeneralResult<SharedMemory> createSharedMemory(size_t size);
86 
87 // Duplicates `fd` and takes ownership of the duplicate.
88 // Precondition: size > 0
89 GeneralResult<SharedMemory> createSharedMemoryFromFd(size_t size, int prot, int fd, size_t offset);
90 
91 #ifdef __ANDROID__
92 // Precondition: ahwb != nullptr
93 GeneralResult<SharedMemory> createSharedMemoryFromAHWB(AHardwareBuffer* ahwb, bool takeOwnership);
94 #endif  // __ANDROID__
95 
96 // Precondition: memory != nullptr
97 size_t getSize(const SharedMemory& memory);
98 
99 bool isAhwbBlob(const Memory::HardwareBuffer& memory);
100 
101 // Precondition: memory != nullptr
102 bool isAhwbBlob(const SharedMemory& memory);
103 
104 struct Mapping {
105     std::variant<const void*, void*> pointer;
106     size_t size;
107     std::any context;
108 };
109 
110 GeneralResult<Mapping> map(const SharedMemory& memory);
111 
112 bool flush(const Mapping& mapping);
113 
114 // Indicates if the object contains no pointer-based data that could be relocated to shared memory.
115 bool hasNoPointerData(const Model& model);
116 bool hasNoPointerData(const Request& request);
117 
118 // Relocate pointer-based data to shared memory. If `model` has no Operand::LifeTime::POINTER data,
119 // the function returns with a reference to `model`. If `model` has Operand::LifeTime::POINTER data,
120 // the model is copied to `maybeModelInSharedOut` with the POINTER data relocated to a memory pool,
121 // and the function returns with a reference to `*maybeModelInSharedOut`.
122 GeneralResult<std::reference_wrapper<const Model>> flushDataFromPointerToShared(
123         const Model* model, std::optional<Model>* maybeModelInSharedOut);
124 
125 // Record a relocation mapping between pointer-based data and shared memory.
126 // Only two specializations of this template may exist:
127 // - RelocationInfo<const void*> for request inputs
128 // - RelocationInfo<void*> for request outputs
129 template <typename PointerType>
130 struct RelocationInfo {
131     PointerType data;
132     size_t length;
133     size_t offset;
134 };
135 using InputRelocationInfo = RelocationInfo<const void*>;
136 using OutputRelocationInfo = RelocationInfo<void*>;
137 
138 // Keep track of the relocation mapping between pointer-based data and shared memory pool,
139 // and provide method to copy the data between pointers and the shared memory pool.
140 // Only two specializations of this template may exist:
141 // - RelocationTracker<InputRelocationInfo> for request inputs
142 // - RelocationTracker<OutputRelocationInfo> for request outputs
143 template <typename RelocationInfoType>
144 class RelocationTracker {
145    public:
create(std::vector<RelocationInfoType> relocationInfos,SharedMemory memory)146     static GeneralResult<std::unique_ptr<RelocationTracker>> create(
147             std::vector<RelocationInfoType> relocationInfos, SharedMemory memory) {
148         auto mapping = NN_TRY(map(memory));
149         return std::make_unique<RelocationTracker<RelocationInfoType>>(
150                 std::move(relocationInfos), std::move(memory), std::move(mapping));
151     }
152 
RelocationTracker(std::vector<RelocationInfoType> relocationInfos,SharedMemory memory,Mapping mapping)153     RelocationTracker(std::vector<RelocationInfoType> relocationInfos, SharedMemory memory,
154                       Mapping mapping)
155         : kRelocationInfos(std::move(relocationInfos)),
156           kMemory(std::move(memory)),
157           kMapping(std::move(mapping)) {}
158 
159     // Specializations defined in CommonUtils.cpp.
160     // For InputRelocationTracker, this method will copy pointer data to the shared memory pool.
161     // For OutputRelocationTracker, this method will copy shared memory data to the pointers.
162     void flush() const;
163 
164    private:
165     const std::vector<RelocationInfoType> kRelocationInfos;
166     const SharedMemory kMemory;
167     const Mapping kMapping;
168 };
169 using InputRelocationTracker = RelocationTracker<InputRelocationInfo>;
170 using OutputRelocationTracker = RelocationTracker<OutputRelocationInfo>;
171 
172 struct RequestRelocation {
173     std::unique_ptr<InputRelocationTracker> input;
174     std::unique_ptr<OutputRelocationTracker> output;
175 };
176 
177 // Relocate pointer-based data to shared memory. If `request` has no
178 // Request::Argument::LifeTime::POINTER data, the function returns with a reference to `request`. If
179 // `request` has Request::Argument::LifeTime::POINTER data, the request is copied to
180 // `maybeRequestInSharedOut` with the POINTER data relocated to a memory pool, and the function
181 // returns with a reference to `*maybeRequestInSharedOut`. The `relocationOut` will be set to track
182 // the input and output relocations.
183 //
184 // Unlike `flushDataFromPointerToShared`, this method will not copy the input pointer data to the
185 // shared memory pool. Use `relocationOut` to flush the input or output data after the call.
186 GeneralResult<std::reference_wrapper<const Request>> convertRequestFromPointerToShared(
187         const Request* request, uint32_t alignment, uint32_t padding,
188         std::optional<Request>* maybeRequestInSharedOut, RequestRelocation* relocationOut);
189 
190 }  // namespace android::nn
191 
192 #endif  // ANDROID_PACKAGES_MODULES_NEURALNETWORKS_COMMON_TYPES_NNAPI_SHARED_MEMORY_H
193