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 #include "BufferTracker.h"
18 
19 #include "HalInterfaces.h"
20 
21 #include <android-base/macros.h>
22 #include <nnapi/TypeUtils.h>
23 
24 #include <memory>
25 #include <mutex>
26 #include <set>
27 #include <stack>
28 #include <utility>
29 #include <vector>
30 
31 namespace android::nn {
32 
create(uint32_t size,std::set<AidlHalPreparedModelRole> roles,const Operand & operand)33 std::shared_ptr<AidlManagedBuffer> AidlManagedBuffer::create(
34         uint32_t size, std::set<AidlHalPreparedModelRole> roles, const Operand& operand) {
35     std::unique_ptr<uint8_t[]> buffer(new (std::nothrow) uint8_t[size]);
36     if (buffer == nullptr) {
37         return nullptr;
38     }
39     if (isExtension(operand.type)) {
40         LOG(ERROR) << "AidlManagedBuffer cannot handle extension operands.";
41         return nullptr;
42     }
43     return std::make_shared<AidlManagedBuffer>(std::move(buffer), size, std::move(roles), operand);
44 }
45 
AidlManagedBuffer(std::unique_ptr<uint8_t[]> buffer,uint32_t size,std::set<AidlHalPreparedModelRole> roles,const Operand & operand)46 AidlManagedBuffer::AidlManagedBuffer(std::unique_ptr<uint8_t[]> buffer, uint32_t size,
47                                      std::set<AidlHalPreparedModelRole> roles,
48                                      const Operand& operand)
49     : kBuffer(std::move(buffer)),
50       kSize(size),
51       kRoles(std::move(roles)),
52       kOperandType(operand.type),
53       kInitialDimensions(operand.dimensions),
54       mUpdatedDimensions(operand.dimensions) {
55     CHECK(!isExtension(kOperandType));
56 }
57 
validateRequest(uint32_t poolIndex,const Request & request,const aidl_hal::IPreparedModel * preparedModel) const58 ErrorStatus AidlManagedBuffer::validateRequest(
59         uint32_t poolIndex, const Request& request,
60         const aidl_hal::IPreparedModel* preparedModel) const {
61     CHECK_LT(poolIndex, request.pools.size());
62     CHECK(std::holds_alternative<Request::MemoryDomainToken>(request.pools[poolIndex]));
63     std::lock_guard<std::mutex> guard(mMutex);
64 
65     bool usedAsInput = false, usedAsOutput = false;
66     for (uint32_t i = 0; i < request.inputs.size(); i++) {
67         if (request.inputs[i].lifetime != Request::Argument::LifeTime::POOL) continue;
68         if (request.inputs[i].location.poolIndex != poolIndex) continue;
69         // Validate if the input role is specified during allocation.
70         if (kRoles.count({preparedModel, IOType::INPUT, i}) == 0) {
71             LOG(ERROR) << "AidlManagedBuffer::validateRequest -- invalid buffer role.";
72             return ErrorStatus::INVALID_ARGUMENT;
73         }
74         if (!mInitialized) {
75             LOG(ERROR)
76                     << "AidlManagedBuffer::validateRequest -- using uninitialized buffer as input "
77                        "request.";
78             return ErrorStatus::GENERAL_FAILURE;
79         }
80         auto combined = combineDimensions(mUpdatedDimensions, request.inputs[i].dimensions);
81         if (!combined.has_value()) {
82             LOG(ERROR) << "AidlManagedBuffer::validateRequest -- incompatible dimensions ("
83                        << toString(mUpdatedDimensions) << " vs "
84                        << toString(request.inputs[i].dimensions) << ")";
85             return ErrorStatus::INVALID_ARGUMENT;
86         }
87         usedAsInput = true;
88     }
89     for (uint32_t i = 0; i < request.outputs.size(); i++) {
90         if (request.outputs[i].lifetime != Request::Argument::LifeTime::POOL) continue;
91         if (request.outputs[i].location.poolIndex != poolIndex) continue;
92         if (usedAsInput || usedAsOutput) {
93             LOG(ERROR) << "AidlManagedBuffer::validateRequest -- using the same device memory for "
94                           "input/output or multiple outputs";
95             return ErrorStatus::INVALID_ARGUMENT;
96         }
97         // Validate if the output role is specified during allocation.
98         if (kRoles.count({preparedModel, IOType::OUTPUT, i}) == 0) {
99             LOG(ERROR) << "AidlManagedBuffer::validateRequest -- invalid buffer role.";
100             return ErrorStatus::INVALID_ARGUMENT;
101         }
102         auto combined = combineDimensions(kInitialDimensions, request.outputs[i].dimensions);
103         if (!combined.has_value()) {
104             LOG(ERROR) << "AidlManagedBuffer::validateRequest -- incompatible dimensions ("
105                        << toString(kInitialDimensions) << " vs "
106                        << toString(request.outputs[i].dimensions) << ")";
107             return ErrorStatus::INVALID_ARGUMENT;
108         }
109         usedAsOutput = true;
110     }
111     return ErrorStatus::NONE;
112 }
113 
validateCopyFrom(const std::vector<uint32_t> & dimensions,uint32_t size) const114 ErrorStatus AidlManagedBuffer::validateCopyFrom(const std::vector<uint32_t>& dimensions,
115                                                 uint32_t size) const {
116     if (size != kSize) {
117         LOG(ERROR) << "AidlManagedBuffer::validateCopyFrom -- invalid memory size: " << kSize
118                    << " vs " << size;
119         return ErrorStatus::INVALID_ARGUMENT;
120     }
121 
122     if (isNonExtensionScalar(kOperandType)) {
123         if (!dimensions.empty()) {
124             LOG(ERROR) << "AidlManagedBuffer::validateCopyFrom -- invalid dimensions for scalar "
125                           "operand: "
126                        << toString(dimensions);
127             return ErrorStatus::INVALID_ARGUMENT;
128         }
129         return ErrorStatus::NONE;
130     }
131 
132     if (dimensions.empty()) {
133         if (tensorHasUnspecifiedDimensions(kOperandType, kInitialDimensions)) {
134             LOG(ERROR) << "AidlManagedBuffer::validateCopyFrom -- the initial dimensions are not "
135                           "fully "
136                           "specified and no dimension update is provided: "
137                        << toString(kInitialDimensions);
138             return ErrorStatus::INVALID_ARGUMENT;
139         }
140     } else {
141         if (tensorHasUnspecifiedDimensions(kOperandType, dimensions)) {
142             LOG(ERROR) << "AidlManagedBuffer::validateCopyFrom -- the updated dimensions are not "
143                           "fully "
144                           "specified: "
145                        << toString(dimensions);
146             return ErrorStatus::INVALID_ARGUMENT;
147         }
148     }
149 
150     const auto combined = combineDimensions(kInitialDimensions, dimensions);
151     if (!combined.has_value()) {
152         LOG(ERROR) << "AidlManagedBuffer::validateCopyFrom -- incompatible dimensions ("
153                    << toString(kInitialDimensions) << " vs " << toString(dimensions) << ")";
154         return ErrorStatus::INVALID_ARGUMENT;
155     }
156     return ErrorStatus::NONE;
157 }
158 
validateCopyTo(uint32_t size) const159 ErrorStatus AidlManagedBuffer::validateCopyTo(uint32_t size) const {
160     if (size != kSize) {
161         LOG(ERROR) << "AidlManagedBuffer::validateCopyTo -- invalid memory size: " << kSize
162                    << " vs " << size;
163         return ErrorStatus::INVALID_ARGUMENT;
164     }
165     std::lock_guard<std::mutex> guard(mMutex);
166     if (!mInitialized) {
167         LOG(ERROR) << "AidlManagedBuffer::validateCopyTo -- using uninitialized buffer as source.";
168         return ErrorStatus::GENERAL_FAILURE;
169     }
170     return ErrorStatus::NONE;
171 }
172 
updateDimensions(const std::vector<uint32_t> & dimensions)173 bool AidlManagedBuffer::updateDimensions(const std::vector<uint32_t>& dimensions) {
174     auto combined = combineDimensions(kInitialDimensions, dimensions);
175     if (!combined.has_value()) {
176         LOG(ERROR) << "AidlManagedBuffer::updateDimensions -- incompatible dimensions ("
177                    << toString(kInitialDimensions) << " vs " << toString(dimensions) << ")";
178         return false;
179     }
180     std::lock_guard<std::mutex> guard(mMutex);
181     mUpdatedDimensions = std::move(combined).value();
182     return true;
183 }
184 
setInitialized(bool initialized)185 void AidlManagedBuffer::setInitialized(bool initialized) {
186     std::lock_guard<std::mutex> guard(mMutex);
187     mInitialized = initialized;
188 }
189 
add(std::shared_ptr<AidlManagedBuffer> buffer)190 std::unique_ptr<AidlBufferTracker::Token> AidlBufferTracker::add(
191         std::shared_ptr<AidlManagedBuffer> buffer) {
192     if (buffer == nullptr) {
193         return nullptr;
194     }
195     std::lock_guard<std::mutex> guard(mMutex);
196     uint32_t token = 0;
197     if (mFreeTokens.empty()) {
198         token = mTokenToBuffers.size();
199         mTokenToBuffers.push_back(std::move(buffer));
200     } else {
201         token = mFreeTokens.top();
202         mFreeTokens.pop();
203         mTokenToBuffers[token] = std::move(buffer);
204     }
205     VLOG(MEMORY) << "AidlBufferTracker::add -- new token = " << token;
206     return std::make_unique<Token>(token, shared_from_this());
207 }
208 
get(uint32_t token) const209 std::shared_ptr<AidlManagedBuffer> AidlBufferTracker::get(uint32_t token) const {
210     std::lock_guard<std::mutex> guard(mMutex);
211     if (mTokenToBuffers.size() <= token || mTokenToBuffers[token] == nullptr) {
212         LOG(ERROR) << "AidlBufferTracker::get -- unknown token " << token;
213         return nullptr;
214     }
215     return mTokenToBuffers[token];
216 }
217 
free(uint32_t token)218 void AidlBufferTracker::free(uint32_t token) {
219     std::lock_guard<std::mutex> guard(mMutex);
220     CHECK_LT(token, mTokenToBuffers.size());
221     CHECK(mTokenToBuffers[token] != nullptr);
222     VLOG(MEMORY) << "AidlBufferTracker::free -- release token = " << token;
223     mTokenToBuffers[token] = nullptr;
224     mFreeTokens.push(token);
225 }
226 
227 }  // namespace android::nn
228