Skip to main content

/

Docs

Log in

Mojo 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 (UnsafePointer[NoneType, 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): 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: UnsafePointer[NoneType, 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1], rank: Int, shape: StaticIntTuple[8], type: DType)

Construct DynamicRankBuffer.

Args:

  • ​data (UnsafePointer[NoneType, 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): 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, -31337, 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[::Int](), create_unknown[::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[::Int](), create_unknown[::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.

Was this page helpful?