[tensor] Support late memory allocation for tensors
authorParichay Kapoor <pk.kapoor@samsung.com>
Mon, 25 Jan 2021 06:31:33 +0000 (15:31 +0900)
committerJijoong Moon <jijoong.moon@samsung.com>
Fri, 5 Feb 2021 01:50:57 +0000 (10:50 +0900)
Support late memory allocation for tensors.
This helps create tensor wrapper elements and allocate memory
later when needed.
This allows finalizing the graph, creating all the tensors and graph
connections for all the layers which can be stored in offline setup
and then loaded directly from file.

This decoupling allows to reuse tensors with every call to train
where we can allocate and deallocate memory without having to create
and do tensor management again and again.

In short term, this is necessary for in-place layer optimization.
See Also #879

The implementation requires caching the source tensor when creating
shared tensors as shared tensors cannot be created unless source
tensor memory has been allocated.

**Self evaluation:**
1. Build test: [x]Passed [ ]Failed [ ]Skipped
2. Run test: [x]Passed [ ]Failed [ ]Skipped

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>
nntrainer/tensor/tensor.cpp
nntrainer/tensor/tensor.h

index feab5ba..ceee9fb 100644 (file)
@@ -89,18 +89,80 @@ static auto rng = [] {
   return rng;
 }();
 
-Tensor::Tensor(const TensorDim &d, const float *buf) : Tensor() {
+Tensor::Tensor(const TensorDim &d, const float *buf, bool alloc_now) : Tensor() {
   if (d.getDataLen() != 0) {
+    if (buf != nullptr && !alloc_now)
+      throw std::invalid_argument("alloc_now must be true to create a tensor with a buffer");
+
     dim = d;
     strides = d.computeStrides();
-    data = std::shared_ptr<float>(new float[d.getDataLen()],
-                                  std::default_delete<float[]>());
+
+    if (alloc_now)
+      data = std::shared_ptr<float>(new float[d.getDataLen()],
+          std::default_delete<float[]>());
 
     if (buf != nullptr)
       copy(buf);
   }
 }
 
+/**
+ * @class SrcSharedTensor
+ * @brief Source of the shared tensor
+ */
+class SrcSharedTensor {
+public:
+  /**
+   * @brief   Constructor for the class
+   */
+  SrcSharedTensor(): src(nullptr), off(0) {}
+
+  SrcSharedTensor(const Tensor *tensor, unsigned int offset): src(tensor), off(offset) {}
+
+  /**
+   * @brief   Get the allocated src tensor
+   */
+  const Tensor *tensor() {
+    if (!src)
+      throw std::runtime_error("Accessing empty src tensor");
+
+    /**
+     * Allocate source tensor if not already allocated.
+     */
+    if (!src->isAllocated())
+      throw std::runtime_error("Source tensor must be allocated before allocating dependent tensors");
+      // src->allocate();
+
+    return src;
+  }
+
+  /**
+   * @brief   Get the offset from the source tensor
+   */
+  unsigned int offset() {
+    return off;
+  }
+
+private:
+  const Tensor *src;   /**< Tensor of the source */
+  unsigned int off;  /**< offset from the source data ptr */
+};
+
+void Tensor::allocate() {
+  if (data)
+    /// already allocated
+    return;
+
+  if (src_tensor) {
+    /// allocate data based on the source tensor
+    data = std::shared_ptr<float>(src_tensor->tensor()->data, src_tensor->tensor()->data.get() + src_tensor->offset());
+  }
+
+  /// allocate new memory for the tensor data
+  data = std::shared_ptr<float>(new float[dim.getDataLen()],
+      std::default_delete<float[]>());
+}
+
 Tensor Tensor::Map(float *buf, unsigned int size, const TensorDim &d,
                    int offset) {
   if (d.getDataLen() == 0 || buf == nullptr) {
@@ -421,9 +483,17 @@ Tensor Tensor::getSharedDataTensor(const TensorDim dim_,
       "Creating shared tensor of size bigger than tensor memory.");
 
   ret.dim = dim_;
-  ret.data = std::shared_ptr<float>(this->data, this->data.get() + offset);
   ret.strides = ret.dim.computeStrides();
 
+  /**
+   * In this case, its the caller's responsibility to ensure that allocate() is
+   * called for this function before operating on this tensor.
+   */
+  if (this->data)
+    ret.data = std::shared_ptr<float>(this->data, this->data.get() + offset);
+  else
+    ret.src_tensor = std::make_shared<SrcSharedTensor>(this, offset);
+
   return ret;
 }
 
index 0e5f037..62beaab 100644 (file)
@@ -43,6 +43,7 @@
 namespace nntrainer {
 
 class LazyTensor;
+class SrcSharedTensor;
 
 /**
  * @class   Tensor Class for Calculation
@@ -57,12 +58,17 @@ public:
     dim(TensorDim()),
     strides{dim.computeStrides()},
     is_contiguous(true),
-    data(nullptr) {}
+    data(nullptr),
+    src_tensor() {}
 
   /**
    * @brief     Constructor of Tensor with dimension/buf
+   * @param d Tensor dim for this tensor
+   * @param bug buffer
+   * @alloc_now If the memory of the tensor must be allocated
+   * @throws if buf is not null and allow_now is false
    */
-  Tensor(const TensorDim &d, const float *buf = nullptr);
+  Tensor(const TensorDim &d, const float *buf = nullptr, bool alloc_now = true);
 
   /**
    * @brief Construct a new Tensor object from a buffer
@@ -190,6 +196,22 @@ public:
   bool operator!=(const Tensor &rhs) const { return !(*this == rhs); }
 
   /**
+   * @brief    Allocate memory for this tensor
+   */
+  void allocate();
+
+  /**
+   * @brief    Deallocate memory for this tensor
+   * @note     This will not necessary free the memory as tensors share memory
+   */
+  void deallocate() { data = nullptr; }
+
+  /**
+   * @brief    Check if the tensor has memory allocated/assigned/associated
+   */
+  bool isAllocated() const { return data == nullptr; }
+
+  /**
    * @brief     return value at specific location
    * @param[in] batch batch location
    * @param[in] c channel location
@@ -900,8 +922,17 @@ private:
   TensorDim dim;
   std::array<unsigned int, MAXDIM> strides;
   bool is_contiguous;
+
   std::shared_ptr<float> data;
 
+  /**<
+   * When using shared_data with tensor, this stores the ptr of the source
+   * tensor which handles the full memory. If tensor data is already allocated,
+   * this does not affect the tensor. If the tensor data is not allocated, and
+   * src_ptr is valid, this tensor will use the memory allocated by the src_ptr
+   */
+  std::shared_ptr<SrcSharedTensor> src_tensor;
+
   template <typename T> void setDist(T dist);
 
   void copy(const float *buf) noexcept;