Merge branch 'rm-buffer' into 'master'

Remove unsed buffer

See merge request !197

mace/core/mace.cc

@@ -14,26 +14,26 @@ ConstTensor::ConstTensor(const std::string &name,
                          const std::vector<int64_t> &dims,
                          const DataType data_type,
                          uint32_t node_id) :
-  name_(name),
-  data_(data),
-  data_size_(std::accumulate(dims.begin(), dims.end(), 1,
-                             std::multiplies<int64_t>())),
-  dims_(dims.begin(), dims.end()),
-  data_type_(data_type),
-  node_id_(node_id) {}
+    name_(name),
+    data_(data),
+    data_size_(std::accumulate(dims.begin(), dims.end(), 1,
+                               std::multiplies<int64_t>())),
+    dims_(dims.begin(), dims.end()),
+    data_type_(data_type),
+    node_id_(node_id) {}
 
 ConstTensor::ConstTensor(const std::string &name,
                          unsigned char *data,
                          const std::vector<int64_t> &dims,
                          const int data_type,
                          uint32_t node_id) :
-  name_(name),
-  data_(data),
-  data_size_(std::accumulate(dims.begin(), dims.end(), 1,
-                             std::multiplies<int64_t>())),
-  dims_(dims.begin(), dims.end()),
-  data_type_(static_cast<DataType>(data_type)),
-  node_id_(node_id) {}
+    name_(name),
+    data_(data),
+    data_size_(std::accumulate(dims.begin(), dims.end(), 1,
+                               std::multiplies<int64_t>())),
+    dims_(dims.begin(), dims.end()),
+    data_type_(static_cast<DataType>(data_type)),
+    node_id_(node_id) {}
 
 const std::string &ConstTensor::name() const {
   return name_;
@@ -151,7 +151,7 @@ void Argument::set_strings(const std::vector<std::string> &value) {
 
 // Node Input
 NodeInput::NodeInput(int node_id, int output_port)
-  : node_id_(node_id), output_port_(output_port) {}
+    : node_id_(node_id), output_port_(output_port) {}
 void NodeInput::CopyFrom(const NodeInput &from) {
   node_id_ = from.node_id();
   output_port_ = from.output_port();
@@ -172,7 +172,7 @@ void NodeInput::set_output_port(int output_port) {
 // OutputShape
 OutputShape::OutputShape() {}
 OutputShape::OutputShape(const std::vector<int64_t> &dims) :
-  dims_(dims.begin(), dims.end()) {}
+    dims_(dims.begin(), dims.end()) {}
 void OutputShape::CopyFrom(const OutputShape &from) {
   auto from_dims = from.dims();
   dims_.resize(from_dims.size());
@@ -359,7 +359,7 @@ void OperatorDef::set_output_type(const std::vector<DataType> &value) {
 
 // MemoryBlock
 MemoryBlock::MemoryBlock(int mem_id, uint32_t x, uint32_t y) :
-  mem_id_(mem_id), x_(x), y_(y) {}
+    mem_id_(mem_id), x_(x), y_(y) {}
 
 int MemoryBlock::mem_id() const {
   return mem_id_;
@@ -511,8 +511,8 @@ const OperatorDef &NetDef::op(const int idx) const {
 
 // Mace Engine
 MaceEngine::MaceEngine(const NetDef *net_def, DeviceType device_type) :
-  op_registry_(new OperatorRegistry()), device_type_(device_type),
-  ws_(new Workspace()), net_(nullptr), hexagon_controller_(nullptr) {
+    op_registry_(new OperatorRegistry()), device_type_(device_type),
+    ws_(new Workspace()), net_(nullptr), hexagon_controller_(nullptr) {
 
   if (device_type == HEXAGON) {
     hexagon_controller_.reset(new HexagonControlWrapper());
@@ -530,9 +530,13 @@ MaceEngine::MaceEngine(const NetDef *net_def, DeviceType device_type) :
     if (!net->Run()) {
       LOG(FATAL) << "Net init run failed";
     }
+    ws_->RemoveUnsedTensor();
     ws_->CreateTensor("mace_input_node:0",
                       GetDeviceAllocator(device_type_),
                       DT_FLOAT);
+    ws_->CreateTensor("mace_output_node:0",
+                      GetDeviceAllocator(device_type_),
+                      DT_FLOAT);
     net_ = std::move(CreateNet(op_registry_, *net_def, ws_.get(), device_type));
   }
 }
@@ -548,14 +552,8 @@ bool MaceEngine::Run(const float *input,
                      const std::vector<index_t> &input_shape,
                      float *output) {
   MACE_CHECK(output != nullptr, "output ptr cannot be NULL");
-  Tensor *input_tensor =
-    ws_->CreateTensor("mace_input_node:0",
-                      GetDeviceAllocator(device_type_),
-                      DT_FLOAT);
-  Tensor *output_tensor =
-    ws_->CreateTensor("mace_output_node:0",
-                      GetDeviceAllocator(device_type_),
-                      DT_FLOAT);
+  Tensor *input_tensor = ws_->GetTensor("mace_input_node:0");
+  Tensor *output_tensor = ws_->GetTensor("mace_output_node:0");
   input_tensor->Resize(input_shape);
   {
     Tensor::MappingGuard input_guard(input_tensor);

mace/core/tensor.h

@@ -70,6 +70,7 @@ class Tensor {
         dtype_(DT_FLOAT),
         buffer_(nullptr),
         data_(nullptr),
+        unused_(false),
         is_image_(false){};
 
   Tensor(Allocator *alloc, DataType type)
@@ -78,6 +79,7 @@ class Tensor {
         dtype_(type),
         buffer_(nullptr),
         data_(nullptr),
+        unused_(false),
         is_image_(false){};
 
   ~Tensor() {
@@ -114,6 +116,8 @@ class Tensor {
 
   inline index_t raw_size() const { return size_ * SizeOfType(); }
 
+  inline const bool unused() const { return unused_; }
+
   inline int64_t NumElements() const {
     return std::accumulate(shape_.begin(), shape_.end(), 1,
                            std::multiplies<int64_t>());
@@ -303,6 +307,10 @@ class Tensor {
     }
   }
 
+  inline void MarkUnused() {
+    this->unused_ = true;
+  }
+
   class MappingGuard {
    public:
     MappingGuard(const Tensor *tensor) : tensor_(tensor) {
@@ -343,6 +351,7 @@ class Tensor {
   mutable void *data_;
   vector<index_t> shape_;
   // Image for opencl
+  bool unused_;
   bool is_image_;
   std::vector<size_t> image_shape_;
 

mace/core/workspace.cc

@@ -47,6 +47,18 @@ const Tensor *Workspace::GetTensor(const string &name) const {
   return nullptr;
 }
 
+
+void Workspace::RemoveUnsedTensor() {
+  auto iter = tensor_map_.begin();
+  auto end_iter = tensor_map_.end();
+  while(iter != end_iter) {
+    auto old_iter = iter++;
+    if(old_iter->second->unused()) {
+      tensor_map_.erase(old_iter);
+    }
+  }
+}
+
 Tensor *Workspace::GetTensor(const string &name) {
   return const_cast<Tensor *>(
       static_cast<const Workspace *>(this)->GetTensor(name));

mace/core/workspace.h

@@ -23,6 +23,8 @@ class Workspace {
 
   bool RemoveTensor(const string &name);
 
+  void RemoveUnsedTensor();
+
   inline bool HasTensor(const string &name) const {
     return tensor_map_.count(name);
   }

mace/kernels/buffer_to_image.h

@@ -27,7 +27,6 @@ struct BufferToImageFunctor : BufferToImageFunctorBase{
                   StatsFuture *future) {
     MACE_NOT_IMPLEMENTED;
   }
-  bool i2b_;
 };
 
 template<typename T>

mace/kernels/opencl/buffer_to_image.cc

@@ -19,6 +19,7 @@ void BufferToImageFunctor<DeviceType::OPENCL, T>::operator()(Tensor *buffer,
   if (!i2b_) {
     CalImage2DShape(buffer->shape(), type, image_shape);
     image->ResizeImage(buffer->shape(), image_shape);
+    buffer->MarkUnused();
   } else {
     image_shape = image->image_shape();
     buffer->Resize(image->shape());