Skip to main content

struct

DynamicRankBuffer

DynamicRankBuffer represents a buffer with unknown rank, shapes and dtype.

It is not as efficient as the statically ranked buffer, but is useful when interacting with external functions. In particular the shape is represented as a fixed (ie _MAX_RANK) array of dimensions to simplify the ABI.

Fields

  • data (DTypePointer[invalid, 0]): The pointer to the buffer.
  • rank (Int): The buffer rank. Has a max value of _MAX_RANK.
  • shape (StaticIntTuple[8]): The dynamic shape of the buffer.
  • type (DType): The dynamic dtype of the buffer.

Implemented traits

AnyType

Methods

__init__

__init__(inout self: Self, /, data: DTypePointer[invalid, 0], rank: Int, shape: StaticIntTuple[8], type: DType)

Construct DynamicRankBuffer.

Args:

  • data (DTypePointer[invalid, 0]): Pointer to the underlying data.
  • rank (Int): Rank of the buffer.
  • shape (StaticIntTuple[8]): Shapes of the buffer.
  • type (DType): dtype of the buffer.

to_buffer

to_buffer[type: DType](self: Self) -> Buffer[$0, Dim(), 0]

Casts DynamicRankBuffer to Buffer.

Parameters:

  • type (DType): dtype of the buffer.

Returns:

Constructed Buffer.

to_ndbuffer

to_ndbuffer[type: DType, rank: Int](self: Self) -> NDBuffer[$0, $1, create_unknown[stdlib::builtin::int::Int](), 0]

Casts the buffer to NDBuffer.

Constraints:

Rank of DynamicRankBuffer must equal rank of NDBuffer.

Parameters:

  • type (DType): dtype of the buffer.
  • rank (Int): Rank of the buffer.

Returns:

Constructed NDBuffer.

to_ndbuffer[type: DType, rank: Int](self: Self, stride: StaticIntTuple[rank]) -> NDBuffer[$0, $1, create_unknown[stdlib::builtin::int::Int](), 0]

Casts the buffer to NDBuffer.

Constraints:

Rank of DynamicRankBuffer must equal rank of NDBuffer.

Parameters:

  • type (DType): dtype of the buffer.
  • rank (Int): Rank of the buffer.

Args:

  • stride (StaticIntTuple[rank]): Strides of the buffer.

Returns:

Constructed NDBuffer.

rank_dispatch

rank_dispatch[func: fn[Int]() capturing -> None](self: Self)

Dispatches the function call based on buffer rank.

Constraints:

Rank must be positive and less or equal to 8.

Parameters:

  • func (fn[Int]() capturing -> None): Function to dispatch. The function should be parametrized on an index parameter, which will be used for rank when the function will be called.

num_elements

num_elements(self: Self) -> Int

Gets number of elements in the buffer.

Returns:

The number of elements in the buffer.

get_shape

get_shape[rank: Int](self: Self) -> StaticIntTuple[$0]

Gets a static tuple representing the buffer shape.

Parameters:

  • rank (Int): Rank of the buffer.

Returns:

A static tuple of size 'Rank' filled with buffer shapes.

dim

dim(self: Self, idx: Int) -> Int

Gets given dimension.

Args:

  • idx (Int): The dimension index.

Returns:

The buffer size on the given dimension.