[manager] Use memory pool for weights

This patch updates the use of memory pool for the weights of the model.
Correspondingly a pool object is added to the manager.
The pool is allocated in the weights allocation.

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>

diff --git a/jni/Android.mk b/jni/Android.mk
index 0203002..9a21c91 100644 (file)
--- a/jni/Android.mk
+++ b/jni/Android.mk
@@ -139,6 +139,8 @@ NNTRAINER_SRCS := $(NNTRAINER_ROOT)/nntrainer/models/neuralnet.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/tensor/var_grad.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/tensor/weight.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/tensor/tensor_dim.cpp \
+                  $(NNTRAINER_ROOT)/nntrainer/tensor/memory_pool.cpp \
+                  $(NNTRAINER_ROOT)/nntrainer/tensor/basic_planner.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/tensor/blas_interface.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/layers/layer_node.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/layers/layer_context.cpp \

diff --git a/nntrainer/tensor/manager.cpp b/nntrainer/tensor/manager.cpp
index 9580967..e5906e4 100644 (file)
--- a/nntrainer/tensor/manager.cpp
+++ b/nntrainer/tensor/manager.cpp
@@ -30,6 +30,7 @@
 #include <vector>
 
 #include <activation_layer.h>
+#include <basic_planner.h>
 #include <flatten_layer.h>
 #include <layer_node.h>
 #include <manager.h>
@@ -249,6 +250,12 @@ Manager::AllocFunc Manager::getAllocFunc(bool is_weight) {
   return allocate_func;
 }
 
+std::pair<unsigned int, unsigned int>
+Manager::getValidity(const std::string &name) {
+  /** @todo calculate validity based on lifespan and usage */
+  return {0, std::numeric_limits<unsigned>::max()};
+}
+
 /**
  * @brief Allocate and initialize the weight variable
  */
@@ -260,9 +267,11 @@ void Manager::initializeWeights() {
   if (LAYER_V2) {
     for (auto &w : weights_v2) {
       w->initializeVariable();
-      // tensor_map(&w->getVariableRef(), requestMemory(w.getDim().size(), 0,
-      // MAX));
+      auto const &t_validity = getValidity(w->getName());
+      tensor_token_map[w->getName()] = pool.requestMemory(
+        w->getVariableRef().bytes(), t_validity.first, t_validity.second);
     }
+    pool.planLayout(BasicPlanner());
   } else {
     if (total_weight_size == 0) {
       ml_logw(
@@ -294,8 +303,10 @@ void Manager::allocateWeights() {
     return;
 
   if (LAYER_V2) {
+    pool.allocate();
     for (auto &w : weights_v2) {
-      w->allocateVariable();
+      w->getVariableRef().setData(
+        pool.getMemory(tensor_token_map[w->getName()]), true);
     }
   } else {
     for (auto &l_w : weights) {
@@ -312,8 +323,11 @@ void Manager::allocateWeights() {
 void Manager::deallocateWeights() {
   if (LAYER_V2) {
     for (auto &w : weights_v2) {
+      /** this just nullifies the set pointer to avoid access to released memory
+       */
       w->deallocateVariable();
     }
+    pool.deallocate();
   } else {
     for (auto &l_w : weights) {
       for (auto &w : l_w) {

diff --git a/nntrainer/tensor/manager.h b/nntrainer/tensor/manager.h
index 06fcd7c..a615082 100644 (file)
--- a/nntrainer/tensor/manager.h
+++ b/nntrainer/tensor/manager.h
@@ -475,6 +475,8 @@ private:
   std::unordered_map<std::string, int>
     name_map; /**< map from output name to its location */
 
+  MemoryPool pool; /**< memory pool for the tensors */
+
   /**< Weights of all the layer in the model to be managed */
   std::vector<std::vector<std::reference_wrapper<Weight>>> weights;
 
@@ -641,6 +643,16 @@ private:
    * @param lifespan The lifespan to be expanded to
    */
   inline void expandLifespan(const std::string &name, TensorLifespan lifespan);
+
+  /**
+   * @brief     Get validity for the given tensor
+   *
+   * @param name Name of the tensor
+   * @return validity for the given tensor
+   * @details the validity will be created using the lifespan and execution
+   * order
+   */
+  std::pair<unsigned int, unsigned int> getValidity(const std::string &name);
 };
 
 } // namespace nntrainer

diff --git a/nntrainer/tensor/memory_pool.cpp b/nntrainer/tensor/memory_pool.cpp
index 8a42d8c..e511990 100644 (file)
--- a/nntrainer/tensor/memory_pool.cpp
+++ b/nntrainer/tensor/memory_pool.cpp
@@ -250,6 +250,24 @@ size_t MemoryPool::calcMinMemoryRequirement() {
                         return val1.second < val2.second;
                       });
   unsigned int last_interval = max_interval.second;
+  /**
+   * as weights stay valid for max duration, ignore this value and get the real
+   * max value
+   */
+  if (last_interval == std::numeric_limits<unsigned int>::max()) {
+    max_interval = *std::max_element(
+      memory_validity.begin(), memory_validity.end(),
+      [last_interval](auto const &val1, auto const &val2) {
+        return ((val2.second != last_interval) && (val1.second < val2.second));
+      });
+    last_interval = max_interval.second;
+    /**
+     * if the second largest number is also numeric_limit, implies that all the
+     * elements are max values. In this case, last_interval is set to 1
+     */
+    if (last_interval == std::numeric_limits<unsigned int>::max())
+      last_interval = 1;
+  }
 
   std::vector<size_t> interval_req(last_interval + 1, 0);
   /**
@@ -260,7 +278,8 @@ size_t MemoryPool::calcMinMemoryRequirement() {
    */
   for (unsigned int idx = 0; idx < memory_size.size(); idx++) {
     for (unsigned int interval = memory_validity[idx].first;
-         interval < memory_validity[idx].second; interval++) {
+         interval < std::min(memory_validity[idx].second, last_interval);
+         interval++) {
       interval_req[interval] += memory_size[idx];
     }
   }

diff --git a/nntrainer/tensor/tensor.h b/nntrainer/tensor/tensor.h
index 7787692..1199a88 100644 (file)
--- a/nntrainer/tensor/tensor.h
+++ b/nntrainer/tensor/tensor.h
@@ -1035,6 +1035,18 @@ public:
    */
   const std::string &getName() const { return name; }
 
+  /**
+   * @brief Set the memory buffer for the tensor
+   *
+   * @param buf the memory buffer
+   * @param init intialize the buffer
+   */
+  void setData(void *buf, bool init = false) {
+    data = std::shared_ptr<float>((float *)buf, [](void *) {});
+    if (init)
+      initialize();
+  }
+
   static constexpr float epsilon = 1e-5;
 
 private: