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 #pragma once 18 19 #include <aidl/android/hardware/neuralnetworks/BnBuffer.h> 20 #include <aidl/android/hardware/neuralnetworks/BnDevice.h> 21 #include <android/binder_auto_utils.h> 22 23 #include <memory> 24 #include <stack> 25 #include <string> 26 #include <utility> 27 #include <vector> 28 29 #include "NeuralNetworksShim.h" 30 #include "SupportLibrary.h" 31 #include "SupportLibraryWrapper.h" 32 33 #include <android-base/macros.h> 34 #include <android-base/thread_annotations.h> 35 36 namespace aidl::android::hardware::neuralnetworks { 37 38 // TODO: This can be templetized. 39 // Keep track of all sl_rapper::Memory and assign each with a unique token. 40 class ShimBufferTracker : public std::enable_shared_from_this<ShimBufferTracker> { 41 DISALLOW_COPY_AND_ASSIGN(ShimBufferTracker); 42 43 public: 44 // A RAII class to help manage the lifetime of the token. 45 // It is only supposed to be constructed in ShimBufferTracker::add. 46 class Token { 47 DISALLOW_COPY_AND_ASSIGN(Token); 48 49 public: Token(uint32_t token,std::shared_ptr<ShimBufferTracker> tracker)50 Token(uint32_t token, std::shared_ptr<ShimBufferTracker> tracker) 51 : kToken(token), kBufferTracker(std::move(tracker)) {} ~Token()52 ~Token() { kBufferTracker->free(kToken); } get()53 uint32_t get() const { return kToken; } 54 55 private: 56 const uint32_t kToken; 57 const std::shared_ptr<ShimBufferTracker> kBufferTracker; 58 }; 59 60 // The factory of ShimBufferTracker. This ensures that the ShimBufferTracker is always managed 61 // by a shared_ptr. create()62 static std::shared_ptr<ShimBufferTracker> create() { 63 return std::make_shared<ShimBufferTracker>(); 64 } 65 66 // Prefer ShimBufferTracker::create. ShimBufferTracker()67 ShimBufferTracker() : mTokenToBuffers(1) {} 68 69 std::unique_ptr<Token> add(std::shared_ptr<::android::nn::sl_wrapper::Memory> buffer); 70 std::shared_ptr<::android::nn::sl_wrapper::Memory> get(uint32_t token) const; 71 72 private: 73 void free(uint32_t token); 74 75 mutable std::mutex mMutex; 76 std::stack<uint32_t, std::vector<uint32_t>> mFreeTokens GUARDED_BY(mMutex); 77 78 // Since the tokens are allocated in a non-sparse way, we use a vector to represent the mapping. 79 // The index of the vector is the token. When the token gets freed, the corresponding entry is 80 // set to nullptr. mTokenToBuffers[0] is always set to nullptr because 0 is an invalid token. 81 std::vector<std::shared_ptr<::android::nn::sl_wrapper::Memory>> mTokenToBuffers 82 GUARDED_BY(mMutex); 83 }; 84 85 } // namespace aidl::android::hardware::neuralnetworks 86