Class DevMemInfo

Simple object that describes a buffer in device memory.

Public Functions

DevMemInfo(void *data, DType dtype, std::shared_ptr<MemoryDescriptor> md, ShapeType shape, ShapeType stride, TensorIndex offset_bytes = 0)

Construct a new DevMemInfo object.

Parameters

• data –

• dtype –

• shape –

• stride –

• offset_bytes –

• stream –

• memory_resource –

DevMemInfo(std::shared_ptr<rmm::device_buffer> buffer, DType dtype, ShapeType shape, ShapeType stride, TensorIndex offset_bytes = 0)

Construct a new DevMemInfo object from an existing rmm::device_buffer.

Parameters

• buffer –

• dtype –

• shape –

• stride –

• offset_bytes –

DevMemInfo(DevMemInfo &&other) = default

TensorIndex bytes() const

Return the number of bytes stored in the underlying buffer.

Returns

TensorIndex

TensorIndex count() const

Return the element count stored in the underlying buffer.

Returns

TensorIndex

TensorIndex offset_bytes() const

Return the number of bytes offset from the head of the buffer.

Returns

TensorIndex

const DType &dtype() const

Return the type of the data stored in the buffer.

Returns

const DType&

TypeId type_id() const

Return the type id of the data stored in the buffer.

Returns

TypeId

const ShapeType &shape() const

Return a reference to the shape.

Returns

const ShapeType&

TensorIndex shape(TensorIndex idx) const

Return a the dimension at idx

Parameters

idx –

Returns

TensorIndex

const ShapeType &stride() const

Return a reference to the stride expressed in elements.

Returns

const ShapeType&

TensorIndex stride(TensorIndex idx) const

Return the stride at idx

Parameters

idx –

Returns

TensorIndex

void *data() const

Returns raw pointer to underlying buffer offset by the offset

Returns

void*

std::shared_ptr<MemoryDescriptor> memory() const

Return the memory descriptor.

Returns

std::shared_ptr<MemoryDescriptor>

std::unique_ptr<rmm::device_buffer> make_new_buffer(TensorIndex bytes) const

Constructs a new rmm buffer with the same stream and memory resource as the current buffer.

Parameters

bytes –

Returns

std::unique_ptr<rmm::device_buffer>