blob.h

#ifndef INCLUDE_UTILS_BLOB_
#define INCLUDE_UTILS_BLOB_
#include <memory>
#include <vector>
#include <glog/logging.h>
#include "proto/model.pb.h"
using std::shared_ptr;
using std::vector;

#define NOT_IMPLEMENTED LOG(FATAL) << "Not implemented function"
inline void MallocHost(void** ptr, size_t size) {
  *ptr = malloc(size);
}

inline void FreeHost(void* ptr) {
  free(ptr);
}

class SyncedMemory {
 public:
  SyncedMemory()
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
        own_cpu_data_(false) {}
  explicit SyncedMemory(size_t size)
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
        own_cpu_data_(false) {}
  ~SyncedMemory();
  const void* cpu_data();
  void set_cpu_data(void* data);
  const void* gpu_data();
  void* mutable_cpu_data();
  void* mutable_gpu_data();
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  SyncedHead head() { return head_; }
  size_t size() { return size_; }

 private:
  void to_cpu();
  void to_gpu();
  void* cpu_ptr_;
  void* gpu_ptr_;
  size_t size_;
  SyncedHead head_;
  bool own_cpu_data_;

};  // class SyncedMemory


template <typename Dtype>
class Blob {
 public:
  Blob(): count_(0), capacity_(0) , version_(-1){}
  Blob(const vector<int>&shape);
  void Reshape(const vector<int>& shape);
  void ReshapeLike(const Blob& other);
  const vector<int>& shape() const{
    return shape_;
  }
  inline int count() const { return count_; }
  void set_version(int v){
    version_=v;
  }
  const int version() const {
    return version_;
  }
  void CopyFrom(const Blob<Dtype>& source, bool reshape = false);

  inline const shared_ptr<SyncedMemory>& data() const {
    CHECK(data_);
    return data_;
  }

  const Dtype* cpu_data() const;
  void set_cpu_data(Dtype* data);
  const Dtype* gpu_data() const;
  Dtype* mutable_cpu_data();
  Dtype* mutable_gpu_data();
  /*
  void FromProto(const BlobProto& proto);
  */
  void ToProto(singa::BlobProto* proto) const;

  Dtype asum_data() const;
  Dtype sum_data() const;

  void ShareData(const Blob& other);
  void Swap(Blob& other);
  shared_ptr<SyncedMemory> data_;
 protected:
  vector<int> shape_;
  int count_;
  int capacity_;
  int version_;
};  // class Blob

#endif // INCLUDE_UTILS_BLOB_