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_BURST_H
18 #define ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_AIDL_UTILS_BURST_H
19 
20 #include <aidl/android/hardware/neuralnetworks/IBurst.h>
21 #include <android-base/scopeguard.h>
22 #include <android-base/thread_annotations.h>
23 #include <nnapi/IBurst.h>
24 #include <nnapi/Result.h>
25 #include <nnapi/Types.h>
26 #include <nnapi/hal/CommonUtils.h>
27 
28 #include <atomic>
29 #include <memory>
30 #include <mutex>
31 #include <optional>
32 #include <unordered_map>
33 #include <utility>
34 
35 // See hardware/interfaces/neuralnetworks/utils/README.md for more information on AIDL interface
36 // lifetimes across processes and for protecting asynchronous calls across AIDL.
37 
38 namespace aidl::android::hardware::neuralnetworks::utils {
39 
40 // Class that adapts aidl_hal::IBurst to nn::IBurst.
41 class Burst final : public nn::IBurst, public std::enable_shared_from_this<Burst> {
42     struct PrivateConstructorTag {};
43 
44   public:
45     /**
46      * Thread-safe, self-cleaning cache that relates an nn::Memory object to a unique int64_t
47      * identifier.
48      */
49     class MemoryCache : public std::enable_shared_from_this<MemoryCache> {
50       public:
51         using Task = std::function<void()>;
52         using Cleanup = ::android::base::ScopeGuard<Task>;
53         using SharedCleanup = std::shared_ptr<const Cleanup>;
54         using WeakCleanup = std::weak_ptr<const Cleanup>;
55 
56         explicit MemoryCache(std::shared_ptr<aidl_hal::IBurst> burst);
57 
58         /**
59          * Get or cache a memory object in the MemoryCache object.
60          *
61          * @param memory Memory object to be cached while the returned `SharedCleanup` is alive.
62          * @return A pair of (1) a unique identifier for the cache entry and (2) a ref-counted
63          *     "hold" object which preserves the cache as long as the hold object is alive.
64          */
65         std::pair<int64_t, SharedCleanup> getOrCacheMemory(const nn::SharedMemory& memory);
66 
67         /**
68          * Get a cached memory object in the MemoryCache object if it exists, otherwise
69          * std::nullopt.
70          *
71          * @param memory Memory object to be cached while the returned `SharedCleanup` is alive.
72          * @return A pair of (1) a unique identifier for the cache entry and (2) a ref-counted
73          *     "hold" object which preserves the cache as long as the hold object is alive. IF the
74          *     cache entry is not present, std::nullopt is returned instead.
75          */
76         std::optional<std::pair<int64_t, SharedCleanup>> getMemoryIfAvailable(
77                 const nn::SharedMemory& memory);
78 
79       private:
80         void tryFreeMemory(const nn::SharedMemory& memory, int64_t identifier);
81 
82         const std::shared_ptr<aidl_hal::IBurst> kBurst;
83         std::mutex mMutex;
84         int64_t mUnusedIdentifier GUARDED_BY(mMutex) = 0;
85         std::unordered_map<nn::SharedMemory, std::pair<int64_t, WeakCleanup>> mCache
86                 GUARDED_BY(mMutex);
87     };
88 
89     static nn::GeneralResult<std::shared_ptr<const Burst>> create(
90             std::shared_ptr<aidl_hal::IBurst> burst);
91 
92     Burst(PrivateConstructorTag tag, std::shared_ptr<aidl_hal::IBurst> burst);
93 
94     // See IBurst::cacheMemory for information.
95     OptionalCacheHold cacheMemory(const nn::SharedMemory& memory) const override;
96 
97     // See IBurst::execute for information.
98     nn::ExecutionResult<std::pair<std::vector<nn::OutputShape>, nn::Timing>> execute(
99             const nn::Request& request, nn::MeasureTiming measure,
100             const nn::OptionalTimePoint& deadline,
101             const nn::OptionalDuration& loopTimeoutDuration) const override;
102 
103     // See IBurst::createReusableExecution for information.
104     nn::GeneralResult<nn::SharedExecution> createReusableExecution(
105             const nn::Request& request, nn::MeasureTiming measure,
106             const nn::OptionalDuration& loopTimeoutDuration) const override;
107 
108     nn::ExecutionResult<std::pair<std::vector<nn::OutputShape>, nn::Timing>> executeInternal(
109             const aidl_hal::Request& request, const std::vector<int64_t>& memoryIdentifierTokens,
110             bool measure, int64_t deadline, int64_t loopTimeoutDuration,
111             const hal::utils::RequestRelocation& relocation) const;
112 
113   private:
114     mutable std::atomic_flag mExecutionInFlight = ATOMIC_FLAG_INIT;
115     const std::shared_ptr<aidl_hal::IBurst> kBurst;
116     const std::shared_ptr<MemoryCache> kMemoryCache;
117 };
118 
119 }  // namespace aidl::android::hardware::neuralnetworks::utils
120 
121 #endif  // ANDROID_HARDWARE_INTERFACES_NEURALNETWORKS_AIDL_UTILS_BURST_H
122