1 | /** |
2 | * Copyright (c) Glow Contributors. See CONTRIBUTORS file. |
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 | #ifndef GLOW_BACKENDS_QUEUEBACKEDDEVICEMANAGER_H |
17 | #define GLOW_BACKENDS_QUEUEBACKEDDEVICEMANAGER_H |
18 | |
19 | #include "glow/Backends/DeviceManager.h" |
20 | #include "glow/Runtime/RuntimeTypes.h" |
21 | #include "glow/Support/ThreadPool.h" |
22 | |
23 | #include <atomic> |
24 | |
25 | namespace glow { |
26 | namespace runtime { |
27 | class QueueBackedDeviceManager : public DeviceManager { |
28 | protected: |
29 | /// Thread which interfaces with the device. |
30 | ThreadPool workThread_; |
31 | |
32 | /// Identifier for next run. |
33 | std::atomic<RunIdentifierTy> nextIdentifier_{1}; |
34 | |
35 | public: |
36 | QueueBackedDeviceManager(const DeviceConfig &config) |
37 | : DeviceManager(config), workThread_(1) {} |
38 | |
39 | virtual ~QueueBackedDeviceManager() { |
40 | ERR_TO_VOID(stop(true)); // will join workThread_ |
41 | } |
42 | |
43 | /// Initialize the device. |
44 | Error init() override { return Error::success(); } |
45 | |
46 | /// Load the provided module into the device, readyCB will be called when |
47 | /// ready to use |
48 | void addNetwork(const Module *module, FunctionMapTy functions, |
49 | ReadyCBTy callback) override { |
50 | workThread_.submit([this, module, f = std::move(functions), |
51 | c = std::move(callback)]() mutable { |
52 | addNetworkImpl(module, std::move(f), std::move(c)); |
53 | }); |
54 | } |
55 | |
56 | /// Remove (and delete) the provided network and all it's functions, freeing |
57 | /// up space on the device. |
58 | void evictNetwork(std::string functionName, |
59 | EvictFunctionCBTy evictCB) override { |
60 | workThread_.submit([this, functionName, evictCB] { |
61 | evictNetworkImpl(functionName, evictCB); |
62 | }); |
63 | } |
64 | |
65 | /// Execute the named Function in an already provided network on the device. |
66 | /// functionName must match the name of a function already added. |
67 | /// The ExecutionContext's PlaceholderBindings should have all Placeholders |
68 | /// allocated. resultCB will be called with the ExecutionContext containing |
69 | /// output tensors filled, and any generated TraceEvents. |
70 | RunIdentifierTy runFunction(std::string functionName, |
71 | std::unique_ptr<ExecutionContext> context, |
72 | ResultCBTy callback) override { |
73 | RunIdentifierTy id = nextIdentifier_++; |
74 | workThread_.submit([this, id, functionName = std::move(functionName), |
75 | context = std::move(context), |
76 | callback = std::move(callback)]() mutable { |
77 | runFunctionImpl(id, std::move(functionName), std::move(context), |
78 | std::move(callback)); |
79 | }); |
80 | return id; |
81 | } |
82 | |
83 | /// Stops execution and shuts down the Device. |
84 | Error stop(bool block = true) override { |
85 | workThread_.stop(block); |
86 | return Error::success(); |
87 | } |
88 | |
89 | protected: |
90 | /// Operator handling methods to be implemented in subclasses (i.e. per Device |
91 | /// type). |
92 | |
93 | /// Load and compile the Module. |
94 | virtual void addNetworkImpl(const Module *, FunctionMapTy, ReadyCBTy) = 0; |
95 | |
96 | /// Remove the module and reclaim its memory. |
97 | virtual void evictNetworkImpl(std::string functionName, |
98 | EvictFunctionCBTy evictCB) = 0; |
99 | |
100 | /// Execute provided Function. |
101 | virtual void runFunctionImpl(RunIdentifierTy, std::string, |
102 | std::unique_ptr<ExecutionContext>, |
103 | ResultCBTy) = 0; |
104 | }; |
105 | } // namespace runtime |
106 | } // namespace glow |
107 | |
108 | #endif // GLOW_BACKENDS_QUEUEBACKEDDEVICEMANAGER_H |
109 | |