LCOV - code coverage report
Current view: top level - nntrainer/tensor - tensor_pool.h (source / functions) Coverage Total Hit
Test: coverage_filtered.info Lines: 66.7 % 27 18
Test Date: 2025-12-14 20:38:17 Functions: 66.7 % 6 4

            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__ */
        

Generated by: LCOV version 2.0-1