[manager] Constraint input and ouput tensor sharing

Add the constraint of the output tensors of a node with the input
tensors of the next nodes. This allows memory pool to understand the
manager that these two inputs/outputs must use the same tensor and
enables tensor sharing.

This is acheieved by requesting already creating tensors with the given
name.

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

diff --git a/nntrainer/graph/network_graph.cpp b/nntrainer/graph/network_graph.cpp
index 5266170..af8282b 100644 (file)
--- a/nntrainer/graph/network_graph.cpp
+++ b/nntrainer/graph/network_graph.cpp
@@ -672,12 +672,21 @@ NetworkGraph::finalizeContext(const std::shared_ptr<LayerNode> &lnode,
                  std::back_inserter(input_dims),
                  [](const Var_Grad *vg) { return vg->getDim(); });
 
+  /** finalize the layer and get the final context */
   auto init_context = lnode->finalize(input_dims);
 
-  std::vector<Var_Grad *> inputs = prev_inputs;
-  if (inputs.empty())
-    inputs =
-      tensor_manager->requestInputs(gnode, init_context.getInputDimensions());
+  /**
+   * Request manager for either a pre-allocated output as input or a newly
+   * allocated input. This is necesary for manager to know when this input node
+   * is going to be used.
+   */
+  std::vector<std::string> input_names;
+  input_names.reserve(prev_inputs.size());
+  std::transform(prev_inputs.begin(), prev_inputs.end(),
+                 std::back_inserter(input_names),
+                 [](auto const &vg) { return vg->getName(); });
+  const std::vector<Var_Grad *> &inputs = tensor_manager->requestInputs(
+    gnode, init_context.getInputDimensions(), input_names);
 
   const std::vector<Var_Grad *> &outputs =
     tensor_manager->requestOutputs(gnode, init_context.getOutputDimensions());
@@ -693,7 +702,10 @@ NetworkGraph::finalizeContext(const std::shared_ptr<LayerNode> &lnode,
 
 int NetworkGraph::initialize() {
   int status = ML_ERROR_NONE;
-  /** this contains the map from name to input tensors for each node */
+  /**
+   * this contains the map from node name to its input tensor names
+   * @note: these input tensors have already been allocated
+   */
   std::unordered_map<std::string, std::vector<Var_Grad *>> input_map;
 
   /** check if the given config of node is of input node */
@@ -701,7 +713,7 @@ int NetworkGraph::initialize() {
     return node->getInputConnections().empty();
   };
 
-  std::vector<Var_Grad *> inputs;
+  std::vector<Var_Grad *> inputs = {};
   for (unsigned int idx = 0; idx < graph.size(); ++idx) {
     auto const &lnode = getSortedLayerNode(idx);
     ml_logd("layer name : %s", lnode->getName().c_str());

diff --git a/nntrainer/tensor/manager.cpp b/nntrainer/tensor/manager.cpp
index 7659153..9580967 100644 (file)
--- a/nntrainer/tensor/manager.cpp
+++ b/nntrainer/tensor/manager.cpp
@@ -781,8 +781,8 @@ Manager::requestWeights(const GraphNode &node,
     tensor_exec_order[gname].push_back(std::get<1>(exec_order));
 
     /** set tensor lifespan */
-    expand_lifespan(vname, TensorLifespan::MAX_LIFESPAN);
-    expand_lifespan(gname, TensorLifespan::BACKWARD_FUNC_LIFESPAN);
+    expandLifespan(vname, TensorLifespan::MAX_LIFESPAN);
+    expandLifespan(gname, TensorLifespan::BACKWARD_FUNC_LIFESPAN);
   }
 
   return ret;
@@ -819,8 +819,8 @@ Manager::requestTensors(const GraphNode &node,
     }
 
     /** set tensor lifespan */
-    expand_lifespan(vname, tspan);
-    expand_lifespan(gname, tspan);
+    expandLifespan(vname, tspan);
+    expandLifespan(gname, tspan);
   }
 
   return ret;
@@ -831,21 +831,42 @@ Manager::requestTensors(const GraphNode &node,
  */
 std::vector<Var_Grad *>
 Manager::requestInputs(const GraphNode &node,
-                       const std::vector<TensorDim> &inputs_dim) {
-  unsigned int count = 0;
+                       const std::vector<TensorDim> &inputs_dim,
+                       const std::vector<std::string> &outputs_name) {
+
   auto const &tspan = TensorLifespan::ITERATION_LIFESPAN;
-  std::vector<Var_Grad::Spec> inputs_spec;
+  std::vector<Var_Grad *> ret;
+
+  if (outputs_name.empty()) {
+    unsigned int count = 0;
+    std::vector<Var_Grad::Spec> inputs_spec;
+
+    std::transform(
+      inputs_dim.begin(), inputs_dim.end(), std::back_inserter(inputs_spec),
+      [&count, &node, &tspan](auto const &elem) {
+        return std::make_tuple(elem, Tensor::Initializer::NONE, true,
+                               node.getName() + std::string(":input") +
+                                 std::to_string(count++),
+                               tspan);
+      });
+
+    ret = requestTensors<Var_Grad>(node, inputs_spec, inputs_v2);
+  } else {
+    ret.reserve(inputs_dim.size());
 
-  std::transform(
-    inputs_dim.begin(), inputs_dim.end(), std::back_inserter(inputs_spec),
-    [&count, &node, &tspan](auto const &elem) {
-      return std::make_tuple(elem, Tensor::Initializer::NONE, true,
-                             node.getName() + std::string(":input") +
-                               std::to_string(count++),
-                             tspan);
-    });
+    /**
+     * Find already allocated output which must match the name and dimensions
+     */
+    for (unsigned int idx = 0; idx < inputs_dim.size(); idx++) {
+      auto output_loc = name_map.at(outputs_name.at(idx));
+      Var_Grad *vg = outputs_v2.at(output_loc).get();
+      if (vg->getDim() != inputs_dim[idx])
+        throw std::invalid_argument(
+          "Dimension mismatch for the requested input");
+      ret.push_back(vg);
+    }
+  }
 
-  auto ret = requestTensors<Var_Grad>(node, inputs_spec, inputs_v2);
   const auto &exec_order = node.getExecutionOrder();
   for (auto const &in : ret) {
     auto const &vname = in->getName();
@@ -859,8 +880,8 @@ Manager::requestInputs(const GraphNode &node,
     tensor_exec_order[gname].push_back(std::get<2>(exec_order));
 
     /** set tensor lifespan */
-    expand_lifespan(vname, tspan);
-    expand_lifespan(gname, tspan);
+    expandLifespan(vname, tspan);
+    expandLifespan(gname, tspan);
   }
 
   return ret;
@@ -905,19 +926,49 @@ Manager::requestOutputs(const GraphNode &node,
     tensor_exec_order[vname].push_back(std::get<2>(exec_order));
 
     /** set tensor lifespan */
-    expand_lifespan(vname, tspan);
-    expand_lifespan(gname, tspan);
+    expandLifespan(vname, tspan);
+    expandLifespan(gname, tspan);
   }
 
   return ret;
 }
 
-void Manager::expand_lifespan(const std::string &name,
-                              TensorLifespan lifespan) {
+void Manager::expandLifespan(const std::string &name, TensorLifespan lifespan) {
   tensor_lifespan_map[name] =
     enum_class_or<TensorLifespan>(tensor_lifespan_map[name], lifespan);
 }
 
+/**
+ * @brief     Create tensors with the given spec
+ */
+std::vector<Var_Grad *> Manager::requestAllocatedOutputsAsInputs(
+  const GraphNode &node, const std::vector<TensorDim> &inputs_dim,
+  const std::vector<std::string> &outputs_name) {
+
+  auto const &tspan = TensorLifespan::ITERATION_LIFESPAN;
+  std::vector<Var_Grad *> ret;
+
+  /** add the execution order and lifespan for the returning tensors */
+  const auto &exec_order = node.getExecutionOrder();
+  for (auto const &in : ret) {
+    auto const &vname = in->getName();
+    auto const &gname = in->getGradientName();
+
+    /** usage for inputs */
+    tensor_exec_order[vname].push_back(std::get<0>(exec_order));
+    tensor_exec_order[vname].push_back(std::get<1>(exec_order));
+
+    /** usage for inputs gradients (outgoing derivatives) */
+    tensor_exec_order[gname].push_back(std::get<2>(exec_order));
+
+    /** set tensor lifespan */
+    expandLifespan(vname, tspan);
+    expandLifespan(gname, tspan);
+  }
+
+  return ret;
+}
+
 std::vector<Weight *> Manager::getWeights() {
   std::vector<Weight *> all_weights;
 

diff --git a/nntrainer/tensor/manager.h b/nntrainer/tensor/manager.h
index 76aa8dd..06fcd7c 100644 (file)
--- a/nntrainer/tensor/manager.h
+++ b/nntrainer/tensor/manager.h
@@ -188,12 +188,20 @@ public:
    *
    * @param node Graph node to extract node identifiers/info
    * @param inputs_dim Specficiation for the tensors
+   * @param outputs_name Name of the already requested output tensors
    *
    * @return created tensors list
+   *
+   * @details create Var_Grads to be used as input of GraphNode with the
+   * inputs_dim as their spec. If the outputs_name is provided, the returned
+   * Var_Grad share tensors with the already allocated Var_Grad for outputs,
+   * named with outputs_name. In this case, the input_dim and the shape of the
+   * output_tensors must match. If the outputs_name are empty, then new tensors
+   * will be allocated.
    */
   std::vector<Var_Grad *>
-  requestInputs(const GraphNode &node,
-                const std::vector<TensorDim> &inputs_dim);
+  requestInputs(const GraphNode &node, const std::vector<TensorDim> &inputs_dim,
+                const std::vector<std::string> &outputs_name = {});
 
   /**
    * @brief     Create tensors with the given spec
@@ -208,6 +216,19 @@ public:
                  const std::vector<TensorDim> &outputs_spec);
 
   /**
+   * @brief     Create tensors with the given spec and name
+   *
+   * @param node Graph node to extract node identifiers/info
+   * @param tensors_dim Specficiation for the tensors
+   *
+   * @return created tensors list
+   */
+  std::vector<Var_Grad *>
+  requestAllocatedOutputsAsInputs(const GraphNode &node,
+                                  const std::vector<TensorDim> &tensors_dim,
+                                  const std::vector<std::string> &outputs_name);
+
+  /**
    * @brief     Get all the weights
    *
    * @return    return all the weights
@@ -451,6 +472,9 @@ private:
   std::unordered_map<std::string, int>
     tensor_token_map; /**< map from tensor to its memory token */
 
+  std::unordered_map<std::string, int>
+    name_map; /**< map from output name to its location */
+
   /**< Weights of all the layer in the model to be managed */
   std::vector<std::vector<std::reference_wrapper<Weight>>> weights;
 
@@ -600,6 +624,7 @@ private:
                                     " with same name");
 
       tensor_exec_order[ts_name] = {};
+      name_map[ts_name] = layer_objs_list.size() - 1;
     }
 
     std::transform(layer_objs_list.begin() + current_size,
@@ -615,7 +640,7 @@ private:
    * @param name The name of the tensor
    * @param lifespan The lifespan to be expanded to
    */
-  inline void expand_lifespan(const std::string &name, TensorLifespan lifespan);
+  inline void expandLifespan(const std::string &name, TensorLifespan lifespan);
 };
 
 } // namespace nntrainer

diff --git a/nntrainer/utils/util_func.h b/nntrainer/utils/util_func.h
index f1a442f..30638f4 100644 (file)
--- a/nntrainer/utils/util_func.h
+++ b/nntrainer/utils/util_func.h
@@ -246,7 +246,7 @@ bool istrequal(const std::string &a, const std::string &b);
  * @param e1 enum value
  * @param e2 enum value
  *
- * @return enum value after performing AND operation
+ * @return enum value after performing logical AND operation
  */
 template <typename T, typename C = int>
 bool enum_class_logical_and(T e1, T e2) {
@@ -262,7 +262,7 @@ bool enum_class_logical_and(T e1, T e2) {
  * @param e1 enum value
  * @param e2 enum value
  *
- * @return enum value after performing AND operation
+ * @return enum value after performing OR operation
  */
 template <typename T, typename C = int> T enum_class_or(T e1, T e2) {
   C i1 = static_cast<int>(e1);