Template Class TensorStorage

Class Documentation

template<class Trait>
class mgb::TensorStorage

manager for raw tensor memory

It contains no dtype information and all sizes are measured in bytes.

Note that ensure_size() is lazy, and memory allocation only happens when ptr() or sub() is called

Public Types

using RawStorage = std::shared_ptr<dt_byte>

Public Functions

TensorStorage() = default
TensorStorage(CompNode comp_node)
TensorStorage(TensorStorage&&) noexcept = default
TensorStorage &operator=(TensorStorage&&) noexcept = default
TensorStorage(const TensorStorage &rhs)
TensorStorage &operator=(const TensorStorage &rhs)
bool valid_span(const TensorLayout::Span &span) const

whether given tensor span is valid in this storage

TensorStorage &ensure_size(size_t sz)

ensure that its space could hold at least sz bytes


  1. This method is lazy; size would only be changed when memory must be accessed.

  2. This method would only grow storage, but it would not release memory

TensorStorage sub(ptrdiff_t offset) const

return a subtensor that shares the memory; the returned subtensor is not allowed to realloc

  • offset: offset given in bytes

dt_byte *ptr() const

apply lazy resize and get ptr

size_t size() const

usable size in bytes until end of allocated block

CompNode comp_node() const

get underlying comp node; error would be raised if it is invalid

CompNode comp_node_allow_invalid() const

get underlying comp node and allow it to be invalid

bool comp_node_valid() const

whether underlying comp_node is valid

bool empty() const

whether this tensor has no valid element (either due to reaching end of mem chunk or no mem allocated)

TensorStorage &comp_node(CompNode node, bool allow_mem_node_change = false)

chain-style computing node setter

note that if allow_mem_node_change is true and memory node is changed, the underlying data would be released and this tensor would become empty

template<class RTrait>
void copy_from(const TensorStorage<RTrait> &src, size_t size) const

copy from another TensorStorage, possibly of other storage type

This storage must have been initialized

  • size: number of bytes to be copied; must not exceed size of this or src

void reset(CompNode node, size_t size, RawStorage data)

reset the tensor storage to given memory area

template<bool x = true, typename = std::enable_if_t<x && std::is_same<Trait, HostTensorStorageTrait>::value>>
DeviceTensorStorage proxy_to_default_cpu() const

make a DeviceTensorStorage on default_cpu that shares memory with this

this must be a HostTensorStorage. Alignment not checked.

size_t use_count() const

shortcut for raw_storage().use_count(), but won’t trigger lazy alloc

bool has_no_real_storage() const

whether current capacity is 0 (so we are waiting for lazy init)

const RawStorage &raw_storage() const

get underlying raw reference-counted storage

Public Static Functions

template<class RTrait, typename = typename std::enable_if<!std::is_same<Trait, RTrait>::value>::type>
TensorStorage make_proxy(const TensorStorage<RTrait> &src)

make a TensorStorage that shares memory with another TensorStorage some different storage type

This method can be used to convert between HostTensorStorage and DeviceTensorStorage; src must be on CPU memory node.