blob.h

/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
* 
*   http://www.apache.org/licenses/LICENSE-2.0
* 
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*
*************************************************************/

#ifndef SINGA_UTILS_BLOB_H_
#define SINGA_UTILS_BLOB_H_

#include <glog/logging.h>
#include <memory>
#include <vector>
#include "proto/common.pb.h"

namespace singa {

inline void MallocHost(void** ptr, size_t size) {
  *ptr = malloc(size);
}

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

class SyncedMemory {
 public:
  enum SyncedHead { UNINITIALIZED,
                    HEAD_AT_CPU,
                    HEAD_AT_GPU,
                    SYNCED };

  SyncedMemory() {}
  explicit SyncedMemory(size_t size) : size_(size) {}
  ~SyncedMemory();

  const void* cpu_data();
  const void* gpu_data();
  void* mutable_cpu_data();
  void* mutable_gpu_data();
  void set_cpu_data(void* data);
  inline SyncedHead head() { return head_; }
  inline size_t size() { return size_; }

 private:
  void to_cpu();
  void to_gpu();

  void* cpu_ptr_ = nullptr;
  void* gpu_ptr_ = nullptr;
  size_t size_ = 0;
  SyncedHead head_ = UNINITIALIZED;
  bool own_cpu_data_ = false;
};  // class SyncedMemory


template <typename Dtype>
class Blob {
 public:
  Blob() {}
  explicit Blob(const std::vector<int>& shape) { Reshape(shape); }
  void Reshape(const std::vector<int>& shape);
  void ReshapeLike(const Blob& other);
  void CopyFrom(const Blob<Dtype>& source);
  void CopyFrom(const Blob<Dtype>& source, bool reshape);
  void FromProto(const singa::BlobProto& proto);
  void ToProto(singa::BlobProto* proto) const;
  void ShareData(const Blob& other);
  void Swap(Blob& other);
  inline const std::vector<int>& shape() const { return shape_; }
  inline int count() const { return count_; }
  inline const int version() const { return version_; }
  inline void set_version(int v) { version_ = v; }
  inline const Dtype* cpu_data() const {
    CHECK(data_);
    return static_cast<const Dtype*>(data_->cpu_data());
  }
  inline void set_cpu_data(Dtype* data) {
    CHECK(data);
    data_->set_cpu_data(data);
  }
  inline const Dtype* gpu_data() const {
    CHECK(data_);
    return static_cast<const Dtype*>(data_->gpu_data());
  }
  inline Dtype* mutable_cpu_data() {
    CHECK(data_);
    return static_cast<Dtype*>(data_->mutable_cpu_data());
  }
  inline Dtype* mutable_gpu_data() {
    CHECK(data_);
    return static_cast<Dtype*>(data_->mutable_gpu_data());
  }
  Dtype asum_data() const;
  Dtype sum_data() const;

 protected:
  std::shared_ptr<SyncedMemory> data_ = nullptr;
  std::vector<int> shape_;
  int count_ = 0;
  int capacity_ = 0;
  int version_ = -1;
};  // class Blob

}  // namespace singa

#endif  // SINGA_UTILS_BLOB_H_