1 /*
2  * Copyright (C) 2021 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_HARDWARE_INTERFACES_NEURALNETWORKS_AIDL_UTILS_BUFFER_TRACKER_H
18 #define ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_AIDL_UTILS_BUFFER_TRACKER_H
19 
20 #include <android-base/macros.h>
21 #include <android-base/thread_annotations.h>
22 
23 #include <map>
24 #include <memory>
25 #include <mutex>
26 #include <set>
27 #include <stack>
28 #include <utility>
29 #include <vector>
30 
31 #include "nnapi/hal/aidl/HalInterfaces.h"
32 #include "nnapi/hal/aidl/ValidateHal.h"
33 
34 namespace android::nn {
35 
36 // This class manages a CPU buffer allocated on heap and provides validation methods.
37 class AidlManagedBuffer {
38   public:
39     static std::shared_ptr<AidlManagedBuffer> create(uint32_t size,
40                                                      std::set<AidlHalPreparedModelRole> roles,
41                                                      const Operand& operand);
42 
43     // Prefer AidlManagedBuffer::create.
44     AidlManagedBuffer(std::unique_ptr<uint8_t[]> buffer, uint32_t size,
45                       std::set<AidlHalPreparedModelRole> roles, const Operand& operand);
46 
getPointer()47     uint8_t* getPointer() const { return kBuffer.get(); }
getSize()48     uint32_t getSize() const { return kSize; }
49 
50     // "poolIndex" is the index of this buffer in the request.pools.
51     ErrorStatus validateRequest(uint32_t poolIndex, const Request& request,
52                                 const aidl_hal::IPreparedModel* preparedModel) const;
53 
54     // "size" is the byte size of the Memory provided to the copyFrom or copyTo method.
55     ErrorStatus validateCopyFrom(const std::vector<uint32_t>& dimensions, uint32_t size) const;
56     ErrorStatus validateCopyTo(uint32_t size) const;
57 
58     bool updateDimensions(const std::vector<uint32_t>& dimensions);
59     void setInitialized(bool initialized);
60 
61   private:
62     mutable std::mutex mMutex;
63     const std::unique_ptr<uint8_t[]> kBuffer;
64     const uint32_t kSize;
65     const std::set<AidlHalPreparedModelRole> kRoles;
66     const OperandType kOperandType;
67     const std::vector<uint32_t> kInitialDimensions;
68     std::vector<uint32_t> mUpdatedDimensions GUARDED_BY(mMutex);
69     bool mInitialized GUARDED_BY(mMutex) = false;
70 };
71 
72 // Keep track of all AidlManagedBuffers and assign each with a unique token.
73 class AidlBufferTracker : public std::enable_shared_from_this<AidlBufferTracker> {
74     DISALLOW_COPY_AND_ASSIGN(AidlBufferTracker);
75 
76   public:
77     // A RAII class to help manage the lifetime of the token.
78     // It is only supposed to be constructed in AidlBufferTracker::add.
79     class Token {
80         DISALLOW_COPY_AND_ASSIGN(Token);
81 
82       public:
Token(uint32_t token,std::shared_ptr<AidlBufferTracker> tracker)83         Token(uint32_t token, std::shared_ptr<AidlBufferTracker> tracker)
84             : kToken(token), kBufferTracker(std::move(tracker)) {}
~Token()85         ~Token() { kBufferTracker->free(kToken); }
get()86         uint32_t get() const { return kToken; }
87 
88       private:
89         const uint32_t kToken;
90         const std::shared_ptr<AidlBufferTracker> kBufferTracker;
91     };
92 
93     // The factory of AidlBufferTracker. This ensures that the AidlBufferTracker is always managed
94     // by a shared_ptr.
create()95     static std::shared_ptr<AidlBufferTracker> create() {
96         return std::make_shared<AidlBufferTracker>();
97     }
98 
99     // Prefer AidlBufferTracker::create.
AidlBufferTracker()100     AidlBufferTracker() : mTokenToBuffers(1) {}
101 
102     std::unique_ptr<Token> add(std::shared_ptr<AidlManagedBuffer> buffer);
103     std::shared_ptr<AidlManagedBuffer> get(uint32_t token) const;
104 
105   private:
106     void free(uint32_t token);
107 
108     mutable std::mutex mMutex;
109     std::stack<uint32_t, std::vector<uint32_t>> mFreeTokens GUARDED_BY(mMutex);
110 
111     // Since the tokens are allocated in a non-sparse way, we use a vector to represent the mapping.
112     // The index of the vector is the token. When the token gets freed, the corresponding entry is
113     // set to nullptr. mTokenToBuffers[0] is always set to nullptr because 0 is an invalid token.
114     std::vector<std::shared_ptr<AidlManagedBuffer>> mTokenToBuffers GUARDED_BY(mMutex);
115 };
116 
117 }  // namespace android::nn
118 
119 #endif  // ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_AIDL_UTILS_BUFFER_TRACKER_H
120