Open ysh329 opened 4 years ago
mace/tensor.h at master · XiaoMi/mace https://github.com/XiaoMi/mace/blob/master/mace/core/tensor.h
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
Mace
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