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
25namespace glow {
26namespace runtime {
27class QueueBackedDeviceManager : public DeviceManager {
28protected:
29 /// Thread which interfaces with the device.
30 ThreadPool workThread_;
31
32 /// Identifier for next run.
33 std::atomic<RunIdentifierTy> nextIdentifier_{1};
34
35public:
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
89protected:
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