yuenshome / yuenshome.github.io

https://yuenshome.github.io
MIT License
81 stars 15 forks source link

Mace、MNN的tensor与buffer #86

Open ysh329 opened 4 years ago

ysh329 commented 4 years ago

Mace

image

Tensor

mace/tensor.h at master · XiaoMi/mace https://github.com/XiaoMi/mace/blob/master/mace/core/tensor.h

Buffer

mace/buffer.h at master · XiaoMi/mace https://github.com/XiaoMi/mace/blob/master/mace/core/buffer.h

namespace mace {
namespace core {
enum BufferType {
  BT_BUFFER,
  BT_IMAGE,
};
}  // namespace core

class BufferBase {
 public:
  BufferBase() : size_(0) {}
  explicit BufferBase(index_t size) : size_(size) {}
  virtual ~BufferBase() {}

  virtual core::BufferType buffer_type() const = 0;
  virtual void *buffer() = 0;
  virtual const void *raw_data() const = 0;
  virtual void *raw_mutable_data() = 0;
  virtual MaceStatus Allocate(index_t nbytes) = 0;
  virtual MaceStatus Allocate(const std::vector<size_t> &shape,
                              DataType data_type) = 0;
  virtual void *Map(index_t offset,
                    index_t length,
                    std::vector<size_t> *pitch) const = 0;
  virtual void UnMap(void *mapped_ptr) const = 0;
  virtual void Map(std::vector<size_t> *pitch) = 0;
  virtual void UnMap() = 0;
  virtual MaceStatus Resize(index_t nbytes) = 0;
  virtual void Copy(void *src, index_t offset, index_t length) = 0;
  virtual bool OnHost() const = 0;
  virtual void Clear() = 0;
  virtual void Clear(index_t size) = 0;
  virtual const std::vector<size_t> shape() const = 0;
  virtual index_t offset() const { return 0; }

  template <typename T>
  const T *data() const {
    return reinterpret_cast<const T *>(raw_data());
  }

  template <typename T>
  T *mutable_data() {
    return reinterpret_cast<T *>(raw_mutable_data());
  }

  index_t size() const { return size_; }

 protected:
  index_t size_;
};

class Buffer : public BufferBase {
 public:
  explicit Buffer(Allocator *allocator)
      : BufferBase(0),
        allocator_(allocator),
        buf_(nullptr),
        mapped_buf_(nullptr),
        is_data_owner_(true) {}

  Buffer(Allocator *allocator, void *data, index_t size)
      : BufferBase(size),
        allocator_(allocator),
        buf_(data),
        mapped_buf_(nullptr),
        is_data_owner_(false) {}

  virtual ~Buffer();
  core::BufferType buffer_type() const;
  void *buffer();
  const void *raw_data() const;
  void *raw_mutable_data();
  MaceStatus Allocate(index_t nbytes);
  MaceStatus Allocate(const std::vector<size_t> &shape,
                      DataType data_type);
  void *Map(index_t offset, index_t length, std::vector<size_t> *pitch) const;
  void UnMap(void *mapped_ptr) cons;
  void Map(std::vector<size_t> *pitch);
  void UnMap();
  MaceStatus Resize(index_t nbytes);
  void Copy(void *src, index_t offset, index_t length);
  bool OnHost() const;
  void Clear();
  void Clear(index_t size);
  const std::vector<size_t> shape() const;
 protected:
  Allocator *allocator_;
  void *buf_;
  void *mapped_buf_;
  bool is_data_owner_;

  MACE_DISABLE_COPY_AND_ASSIGN(Buffer);
};

class Image : public BufferBase {
 public:
  explicit Image(Allocator *allocator)
      : BufferBase(0),
        allocator_(allocator),
        buf_(nullptr),
        mapped_buf_(nullptr) {}

  virtual ~Image() {
    if (mapped_buf_ != nullptr) {
      UnMap();
    }
    if (buf_ != nullptr) {
      allocator_->DeleteImage(buf_);
    }
  }

  inline DataType dtype() const;
  core::BufferType buffer_type() const;
  void *buffer();
  const void *raw_data() const;
  void *raw_mutable_data();
  MaceStatus Allocate(index_t nbytes);
  MaceStatus Allocate(const std::vector<size_t> &shape,
                      DataType data_type);
  void *Map(index_t offset, index_t length, std::vector<size_t> *pitch) const;
  void UnMap(void *mapped_ptr) const;
  void Map(std::vector<size_t> *pitch);
  void UnMap();
  MaceStatus Resize(index_t size);
  void Copy(void *src, index_t offset, index_t length);
  bool OnHost() const { return allocator_->OnHost(); }
  void Clear();
  void Clear(index_t size);
  const std::vector<size_t> shape() const;

 private:
  Allocator *allocator_;
  std::vector<size_t> shape_;
  DataType data_type_;
  void *buf_;
  void *mapped_buf_;

  MACE_DISABLE_COPY_AND_ASSIGN(Image);
};

class BufferSlice : public BufferBase {
 public:
  BufferSlice()
      : BufferBase(0), buffer_(nullptr), mapped_buf_(nullptr), offset_(0) {}
  BufferSlice(BufferBase *buffer, index_t offset, index_t length)
    : BufferBase(length),
      buffer_(buffer),
      mapped_buf_(nullptr),
      offset_(offset) {
    MACE_CHECK(offset >= 0, "buffer slice offset should >= 0");
    MACE_CHECK(offset + length <= buffer->size(),
               "buffer slice offset + length (",
               offset,
               " + ",
               length,
               ") should <= ",
               buffer->size());
  }
  BufferSlice(const BufferSlice &other)
      : BufferSlice(other.buffer_, other.offset_, other.size_) {}

  virtual ~BufferSlice();
  core::BufferType buffer_type() const;
  void *buffer();
  const void *raw_data() const;
  void *raw_mutable_data();
  MaceStatus Allocate(index_t size);
  MaceStatus Allocate(const std::vector<size_t> &shape,
                      DataType data_type);
  void *Map(index_t offset, index_t length, std::vector<size_t> *pitch) const;
  void UnMap(void *mapped_ptr) const;
  void Map(std::vector<size_t> *pitch);
  void UnMap();
  MaceStatus Resize(index_t size);
  void Copy(void *src, index_t offset, index_t length);
  index_t offset() const;
  bool OnHost() const;
  void Clear();
  void Clear(index_t size);
  const std::vector<size_t> shape() const;

 private:
  BufferBase *buffer_;
  void *mapped_buf_;
  index_t offset_;
};

class ScratchBuffer: public Buffer {
 public:
  explicit ScratchBuffer(Allocator *allocator)
    : Buffer(allocator),
      offset_(0) {}

  ScratchBuffer(Allocator *allocator, void *data, index_t size)
    : Buffer(allocator, data, size),
      offset_(0) {}

  virtual ~ScratchBuffer();
  MaceStatus GrowSize(const index_t size);
  BufferSlice Scratch(index_t size);
  void Rewind(index_t offset = 0);
  index_t offset() const;

 private:
  index_t offset_;
};

}  // namespace mace
ysh329 commented 4 years ago

MNN

Buffer

Tensor