You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
121 lines
4.4 KiB
121 lines
4.4 KiB
/*
|
|
* Copyright (C) 2020 The Android Open Source Project
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#ifndef ANDROID_FRAMEWORKS_ML_NN_COMMON_BUFFER_TRACKER_H
|
|
#define ANDROID_FRAMEWORKS_ML_NN_COMMON_BUFFER_TRACKER_H
|
|
|
|
#include <android-base/macros.h>
|
|
#include <android-base/thread_annotations.h>
|
|
|
|
#include <map>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <set>
|
|
#include <stack>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "CpuExecutor.h"
|
|
#include "LegacyUtils.h"
|
|
#include "nnapi/Types.h"
|
|
#include "nnapi/Validation.h"
|
|
|
|
namespace android::nn {
|
|
|
|
// This class manages a CPU buffer allocated on heap and provides validation methods.
|
|
class ManagedBuffer {
|
|
public:
|
|
static std::shared_ptr<ManagedBuffer> create(uint32_t size, std::set<PreparedModelRole> roles,
|
|
const Operand& operand);
|
|
|
|
// Prefer ManagedBuffer::create.
|
|
ManagedBuffer(std::unique_ptr<uint8_t[]> buffer, uint32_t size,
|
|
std::set<PreparedModelRole> roles, const Operand& operand);
|
|
|
|
RunTimePoolInfo createRunTimePoolInfo() const {
|
|
return RunTimePoolInfo::createFromExistingBuffer(kBuffer.get(), kSize);
|
|
}
|
|
|
|
// "poolIndex" is the index of this buffer in the request.pools.
|
|
ErrorStatus validateRequest(uint32_t poolIndex, const Request& request,
|
|
const IPreparedModel* preparedModel) const;
|
|
|
|
// "size" is the byte size of the Memory provided to the copyFrom or copyTo method.
|
|
ErrorStatus validateCopyFrom(const Dimensions& dimensions, uint32_t size) const;
|
|
ErrorStatus validateCopyTo(uint32_t size) const;
|
|
|
|
bool updateDimensions(const Dimensions& dimensions);
|
|
void setInitialized(bool initialized);
|
|
|
|
private:
|
|
mutable std::mutex mMutex;
|
|
const std::unique_ptr<uint8_t[]> kBuffer;
|
|
const uint32_t kSize;
|
|
const std::set<PreparedModelRole> kRoles;
|
|
const OperandType kOperandType;
|
|
const Dimensions kInitialDimensions;
|
|
Dimensions mUpdatedDimensions GUARDED_BY(mMutex);
|
|
bool mInitialized GUARDED_BY(mMutex) = false;
|
|
};
|
|
|
|
// Keep track of all ManagedBuffers and assign each with a unique token.
|
|
class BufferTracker : public std::enable_shared_from_this<BufferTracker> {
|
|
DISALLOW_COPY_AND_ASSIGN(BufferTracker);
|
|
|
|
public:
|
|
// A RAII class to help manage the lifetime of the token.
|
|
// It is only supposed to be constructed in BufferTracker::add.
|
|
class Token {
|
|
DISALLOW_COPY_AND_ASSIGN(Token);
|
|
|
|
public:
|
|
Token(Request::MemoryDomainToken token, std::shared_ptr<BufferTracker> tracker)
|
|
: kToken(token), kBufferTracker(std::move(tracker)) {}
|
|
~Token() { kBufferTracker->free(kToken); }
|
|
Request::MemoryDomainToken get() const { return kToken; }
|
|
|
|
private:
|
|
const Request::MemoryDomainToken kToken;
|
|
const std::shared_ptr<BufferTracker> kBufferTracker;
|
|
};
|
|
|
|
// The factory of BufferTracker. This ensures that the BufferTracker is always managed by a
|
|
// shared_ptr.
|
|
static std::shared_ptr<BufferTracker> create() { return std::make_shared<BufferTracker>(); }
|
|
|
|
// Prefer BufferTracker::create.
|
|
BufferTracker();
|
|
|
|
std::unique_ptr<Token> add(std::shared_ptr<ManagedBuffer> buffer);
|
|
std::shared_ptr<ManagedBuffer> get(Request::MemoryDomainToken token) const;
|
|
|
|
private:
|
|
void free(Request::MemoryDomainToken token);
|
|
|
|
mutable std::mutex mMutex;
|
|
std::stack<Request::MemoryDomainToken, std::vector<Request::MemoryDomainToken>> mFreeTokens
|
|
GUARDED_BY(mMutex);
|
|
|
|
// Since the tokens are allocated in a non-sparse way, we use a vector to represent the mapping.
|
|
// The index of the vector is the token. When the token gets freed, the corresponding entry is
|
|
// set to nullptr. mTokenToBuffers[0] is always set to nullptr because 0 is an invalid token.
|
|
std::vector<std::shared_ptr<ManagedBuffer>> mTokenToBuffers GUARDED_BY(mMutex);
|
|
};
|
|
|
|
} // namespace android::nn
|
|
|
|
#endif // ANDROID_FRAMEWORKS_ML_NN_COMMON_BUFFER_TRACKER_H
|