Line data Source code
1 : // SPDX-License-Identifier: Apache-2.0
2 : /**
3 : * Copyright (C) 2021 Parichay Kapoor <pk.kapoor@samsung.com>
4 : *
5 : * @file tensor_pool.h
6 : * @date 18 Aug 2021
7 : * @brief This is TensorPool for all requested tensors
8 : * @see https://github.com/nnstreamer/nntrainer
9 : * @author Parichay Kapoor <pk.kapoor@samsung.com>
10 : * @author Jihoon Lee <jhoon.it.lee@samsung.com>
11 : * @bug No known bugs except for NYI items
12 : *
13 : */
14 :
15 : #ifndef __TENSOR_POOL_H__
16 : #define __TENSOR_POOL_H__
17 : #ifdef __cplusplus
18 :
19 : #include <functional>
20 : #include <limits>
21 : #include <memory>
22 : #include <unordered_map>
23 : #include <variant>
24 : #include <vector>
25 :
26 : #include <cache_loader.h>
27 : #include <cache_pool.h>
28 : #include <common.h>
29 : #include <tensor.h>
30 : #include <tensor_wrap_specs.h>
31 :
32 : namespace nntrainer {
33 :
34 : /**
35 : * @class TensorPool
36 : * @brief tensor pool of nntrainer
37 : */
38 : class TensorPool {
39 :
40 : public:
41 : static constexpr unsigned PERSIST_END_ORDER =
42 : std::numeric_limits<unsigned>::max();
43 : /**
44 : * @brief Constructor of TensorPool
45 : */
46 1746 : TensorPool() :
47 3492 : mem_pool(std::make_unique<MemoryPool>()), cache_loader(nullptr) {}
48 :
49 : /**
50 : * @brief Constructor of TensorPool
51 : */
52 1388 : TensorPool(
53 : bool enable_fsu, const std::string &fsu_path = "",
54 : const std::string &fsu_name = "",
55 1388 : ml::train::ExecutionMode execution_mode = ml::train::ExecutionMode::TRAIN) {
56 1388 : if (enable_fsu) {
57 : auto cache_pool =
58 : std::make_shared<CachePool>(fsu_path, fsu_name, execution_mode);
59 0 : cache_loader = std::make_unique<CacheLoader>(cache_pool);
60 : mem_pool = cache_pool;
61 : } else {
62 1388 : mem_pool = std::make_shared<MemoryPool>();
63 : }
64 1388 : }
65 :
66 : /**
67 : * @brief Destructor of TensorPool
68 : */
69 6268 : ~TensorPool() = default;
70 :
71 : /**
72 : * @brief reinitialize TensorPool
73 : */
74 1 : void reinitialize() {
75 : name_map.clear();
76 1 : mem_pool = std::make_shared<MemoryPool>();
77 1 : }
78 :
79 : /**
80 : * @brief finalize the requested tensors
81 : * @param planner planner to layout the tensor memories
82 : * @param start_order start value for the order_exec (inclusive)
83 : * @param end_order end value for the order_exec (inclusive)
84 : *
85 : * @details finalize the requested tensors, request memory for them and plan
86 : * layout for their allocations.
87 : */
88 : void finalize(const MemoryPlanner &planner, unsigned int start_order,
89 : unsigned int end_order);
90 :
91 : /**
92 : * @brief Set the batch size for the inputs/outputs of the layers
93 : */
94 : void setBatchSize(const std::string &name, unsigned int batch);
95 :
96 : /**
97 : * @brief Allocate memory for all the managed tensors
98 : */
99 : void allocate(bool init = true);
100 :
101 : /**
102 : * @brief Deallocate memory for all the managed tensors
103 : */
104 : void deallocate();
105 :
106 : /**
107 : * @brief Get execution order for the given tensor
108 : *
109 : * @return The execution order of the tensor
110 : */
111 : const std::vector<unsigned int> &getExecutionOrder(const std::string &name);
112 :
113 : /**
114 : * @brief Get the maximum real memory requirement
115 : *
116 : * @return The real memory requirement with this strategy in bytes
117 : */
118 : size_t size() { return mem_pool->size(); }
119 :
120 : /**
121 : * @brief Get the minimum theoretical memory requirement
122 : *
123 : * @return The theoretical memory requirement with this strategy in bytes
124 : */
125 1317 : size_t minMemoryRequirement() { return mem_pool->minMemoryRequirement(); }
126 :
127 : /**
128 : * @brief Is the tensor pool allocated
129 : *
130 : * @return true if the tensors are allocated, else false
131 : */
132 2418 : bool isAllocated() const { return mem_pool->isAllocated(); }
133 :
134 : /**
135 : * @brief Get the tensor of the given name
136 : *
137 : * @return ptr to the tensor with the given
138 : * @throws if no tensor is found with the given name
139 : */
140 : Tensor *getTensor(const std::string &name) {
141 3913 : return pool[name_map.at(name)].tensor.get();
142 : }
143 :
144 : /**
145 : * @brief Update externally dependent tensors
146 : *
147 : * @param name Name of the tensor
148 : * @param t External tensor
149 : *
150 : * @note Update externally dependent tensors data ptrs from their parents
151 : */
152 : void fillPlaceholder(const std::string &name, const Tensor &t);
153 :
154 : /**
155 : * @brief request placeholder which will be not managed by this tensor pool
156 : * but will be managed externally
157 : *
158 : * @param name Name of the tensor
159 : * @param dim Tensor dimension
160 : * @return Tensor* ptr to the tensor
161 : *
162 : * @note returns empty tensor which must be filled by the caller before use.
163 : */
164 : Tensor *placeholder(const std::string &name, const TensorDim &dim);
165 :
166 : /**
167 : * @brief create a new tensor with the given spec.
168 : *
169 : * @param name Name of this tensor.
170 : * @param dim Tensor dimension.
171 : * @param exec_order The execution orders for this tensor.
172 : * @param lifespan Lifespan of this tensor.
173 : * @param init Initializer of the tensor.
174 : * @param is_weight_grad Identification of weight gradient
175 : *
176 : * @return ptr to the created tensor
177 : *
178 : * @note returns empty tensor which will be filled when allocate is called.
179 : * @note we assume that the caller checks if the exec_order and lifespan are
180 : * compatible.
181 : */
182 : Tensor *request(const std::string &name, const TensorDim &dim,
183 : const std::vector<unsigned int> &exec_order,
184 : TensorLifespan lifespan,
185 3314 : const Initializer &init = Initializer::NONE,
186 : bool is_weight_grad = false);
187 :
188 : /**
189 : * @brief Request tensor which is a view of already requested with the
190 : * given spec
191 : *
192 : * @param name Name of this tensor
193 : * @param reference Name of the reference tensor
194 : * @param dim Tensor dimensions
195 : * @param exec_order The execution orders for this tensors
196 : * @param lifespan Lifespan of this tensor
197 : * @param offset offset from the reference
198 : *
199 : * @return ptr to a tensor which is sharing the same data with
200 : * reference.
201 : *
202 : * @note returns a view tensor which will be filled when the source tensor is
203 : * allocated.
204 : * @note we assume that the caller checks if the exec_order and lifespan are
205 : * compatible.
206 : *
207 : */
208 : Tensor *view(const std::string &name, const std::string &reference,
209 : const TensorDim &dim,
210 : const std::vector<unsigned int> &exec_order,
211 : TensorLifespan lifespan, const size_t offset = 0);
212 :
213 : /**
214 : * @brief extend a tensor life as tensor is being shared.
215 : *
216 : * @param name name of the tensor to extend
217 : * @param dim dimension of the tensor
218 : * @param exec_order exec_order to extend
219 : * @param lifespan extended life span
220 : * @return Tensor* Tensor* the exact tensor which is being extended.
221 : * @note we assume that the caller checks if the exec_order and lifespan are
222 : * compatible.
223 : */
224 : Tensor *extend(const std::string &name, const TensorDim &dim,
225 : const std::vector<unsigned int> &exec_order,
226 : TensorLifespan lifespan);
227 :
228 : /**
229 : * @brief create a new tensor if tensor does not exist else return the tensor
230 : * while extending the tensor's life according to the given arguments.
231 : * @note Created (or extended) tensor is considered identical and managed. It
232 : * is invalid to create a tensor with lifespan::UNMANAGED or dimension and
233 : * initializer is different upon extension.
234 : *
235 : * @param name Name of the tensor
236 : * @param dim dimension
237 : * @param exec_order exec order
238 : * @param lifespan tensor life span
239 : * @param init tensor initializer
240 : * @return Tensor* ptr to either to the existing tensor or newly created
241 : * tensor
242 : */
243 : Tensor *requestOrExtend(const std::string &name, const TensorDim &dim,
244 : const std::vector<unsigned int> &exec_order,
245 : TensorLifespan lifespan,
246 9 : const Initializer &init = Initializer::NONE);
247 :
248 : /**
249 : * @brief reidentify the source of already created tensor (or view).
250 : * @note if @a dest tensor is a view of another tensor, the old source tensor
251 : * of the view will become a view of @a new_src.
252 : *
253 : * @throws std::invalid_argument 1. if the data size required from the
254 : * original source tensor is bigger than the new_src + offset. 2. if new_src
255 : * is a view. Second restriction can be removed, if this is considered as a
256 : * safe behavior.
257 : *
258 : * @param dest identifier for the dest tensor
259 : * @param new_src identifier for the new source tensor
260 : * @param offset offset
261 : */
262 : void reidentifySource(const std::string &dest, const std::string &new_src,
263 : unsigned int offset);
264 :
265 : /**
266 : * @brief flush cache data
267 : *
268 : */
269 : void flushCache();
270 :
271 : /**
272 : * @brief flush cache data except order
273 : *
274 : * @param order except execution order
275 : *
276 : */
277 : void flushCacheExcept(unsigned int order);
278 :
279 : /**
280 : * @brief load cache data by execution order
281 : *
282 : * @param order execution order
283 : */
284 : void loadCacheExec(unsigned int order);
285 :
286 : /**
287 : * @brief load cache data by execution order
288 : *
289 : * @param order execution order
290 : * @return async task id
291 : */
292 : int loadCacheExecAsync(unsigned int order,
293 : TaskExecutor::CompleteCallback complete_callback);
294 :
295 : /**
296 : * @brief check if tensors are loaded for the given execution order.
297 : *
298 : * @param order target execution order
299 : * @return bool true if tensors are loaded, false otherwise.
300 : */
301 : bool checkLoadComplete(unsigned int order);
302 :
303 : /**
304 : * @brief load cache data by execution order
305 : *
306 : * @param order execution order
307 : * @return async task id
308 : */
309 : int flushCacheExecAsync(unsigned int order,
310 : TaskExecutor::CompleteCallback complete_callback);
311 :
312 : /**
313 : * @brief load cache data by execution order
314 : *
315 : * @param id async task id
316 : */
317 : void loadCacheCancel(int id);
318 :
319 : /**
320 : * @brief This function will reset Actives at the given order.
321 : *
322 : */
323 : unsigned int inActive(unsigned int order);
324 :
325 : /**
326 : * @brief set FSU weight path
327 : *
328 : * @param path FSU weight file path
329 : */
330 0 : void setFsuWeightPath(std::string path) {
331 0 : if (mem_pool) {
332 0 : mem_pool->setFsuWeightPath(path);
333 : }
334 0 : }
335 :
336 : /**
337 : * @brief set weight file offset for FSU loading
338 : *
339 : * @param offsets weight file offset
340 : */
341 0 : void setWeightOffset(std::vector<std::pair<size_t, size_t>> offsets) {
342 0 : if (mem_pool) {
343 0 : mem_pool->setWeightOffset(offsets);
344 : }
345 0 : }
346 :
347 : private:
348 : /**
349 : * @brief Source tensor detailed specification
350 : *
351 : */
352 177958 : struct SourceDetails {
353 : unsigned int token; /**< memory token */
354 : TensorLifespan lifespan; /**< life span of the tensor */
355 : std::vector<unsigned int> exec_order; /**< exec order */
356 : std::vector<unsigned int>
357 : dependents; /**< list of dependents to the source */
358 : };
359 :
360 : /**
361 : * @brief Dependent tensor detaild specification
362 : *
363 : */
364 : struct DependentDetails {
365 : unsigned int parent_idx; /**< index to the parent */
366 : unsigned int offset; /**< elementwise offset */
367 : };
368 :
369 : /**
370 : * @brief Spec for storing each request of tensor from tensor pool
371 : * @todo move tensor initialization from tensor class to RequestSpec
372 : */
373 34132 : struct RequestSpec {
374 : bool is_weight_grad; /**< identification of weight gradient */
375 : std::unique_ptr<Tensor> tensor; /**< tensor object itself */
376 : std::variant<SourceDetails, DependentDetails>
377 : details; /**< additional information by its kind */
378 : };
379 :
380 : /**
381 : * @brief check if a tensor exist with the given identifier
382 : *
383 : * @param name name name to check
384 : * @retval true if exist
385 : * @retval false if do not exist
386 : */
387 : bool tensorExist(const std::string &name);
388 :
389 : /**
390 : * @brief Get the view of source Spec from the name
391 : *
392 : * @param name name to get source spec
393 : * @return RequestSpec spec
394 : */
395 : RequestSpec &getSourceSpec(const std::string &name);
396 :
397 : /**
398 : * @brief Expand the lifespan of the tensor with the given name
399 : *
400 : * @param name The name of the tensor
401 : * @param exec_order The execution orders
402 : * @param lifespan The lifespan to be expanded to
403 : * @return source spec for the name
404 : */
405 : RequestSpec &expandLifespan(const std::string &name,
406 : const std::vector<unsigned int> &exec_order,
407 : TensorLifespan lifespan);
408 :
409 : /**
410 : * @brief expand life span with execution time
411 : *
412 : * @param spec specification
413 : * @param exec_order exec order
414 : * @param lifespan life span
415 : */
416 : void expandLifespan(RequestSpec &spec,
417 : const std::vector<unsigned int> &exec_order,
418 : TensorLifespan lifespan);
419 :
420 : /**
421 : * @brief sync dependent tensors from updated source tensor
422 : * @note syncing starting from dependents of dependents is invalid and will
423 : * throw.
424 : *
425 : * @param spec spec with source details to refer to.
426 : */
427 : void syncDependents(const RequestSpec &spec);
428 :
429 : /**
430 : * @brief register a spec after creation
431 : *
432 : * @param spec spec to register
433 : */
434 : Tensor *registerRequestSpec(RequestSpec &&spec);
435 :
436 : /**
437 : * note: unordered_map is not directly used for pool to ensure initialization
438 : * of weights
439 : */
440 : std::vector<RequestSpec> pool; /**< list of requested tensors */
441 : std::unordered_map<std::string, unsigned int>
442 : name_map; /**< indexing of requested tensors */
443 : std::shared_ptr<MemoryPool> mem_pool; /**< memory pool for the tensors */
444 : std::unique_ptr<CacheLoader> cache_loader; /**< memory pool for the tensors */
445 :
446 : /**
447 : * @brief Check if the lifespan leads to long term valitidy
448 : *
449 : * @param lifespan Lifespan for the tensor
450 : *
451 : * @return true if the tensor should be valid for long term, else false
452 : */
453 : bool isTensorLongTerm(const TensorLifespan &lifespan);
454 : };
455 :
456 : } // namespace nntrainer
457 :
458 : #endif /* __cplusplus */
459 : #endif /* __TENSOR_POOL_H__ */
|