Skip to main content
Log in

Mojo struct

NDBuffer

@register_passable(trivial) struct NDBuffer[type: DType, rank: Int, shape: DimList = create_unknown[::Int](), strides: DimList = create_unknown[::Int](), *, alignment: Int = 1, address_space: AddressSpace = 0, exclusive: Bool = True]

An N-dimensional Buffer.

NDBuffer can be parametrized on rank, static dimensions and Dtype. It does not own its underlying pointer.

Parameters

  • type (DType): The element type of the buffer.
  • rank (Int): The rank of the buffer.
  • shape (DimList): The static size (if known) of the buffer.
  • strides (DimList): The strides (if known) of the buffer.
  • alignment (Int): The preferred address alignment of the buffer.
  • address_space (AddressSpace): The address space of the buffer.
  • exclusive (Bool): The underlying memory allocation of the tensor is known only to be accessible through this pointer.

Fields

  • data (UnsafePointer[SIMD[type, 1], address_space=address_space]): The underlying data for the buffer. The pointer is not owned by the NDBuffer.
  • dynamic_shape (IndexList[rank, unsigned=True]): The dynamic value of the shape.
  • dynamic_stride (IndexList[rank, unsigned=True]): The dynamic stride of the buffer.

Implemented traits

AnyType, Copyable, ExplicitlyCopyable, Movable, Sized, Stringable, UnknownDestructibility, Writable

Methods

__init__

__init__() -> Self

Default initializer for NDBuffer. By default the fields are all initialized to 0.

@implicit __init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]) -> Self

Constructs an NDBuffer with statically known rank, shapes and type.

Constraints:

The rank, shapes, and type are known.

Args:

  • ptr (UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): Pointer to the data.

@implicit __init__(other: NDBuffer[type, rank, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]) -> Self

Converts NDBuffers between different variants which do not effect the underlying memory representation.

E.g. this allows implicit conversion between

NDBuffer[type, rank, DimList(1, 2, 3), DimList(6, 6, 1), alignment=16] to NDBuffer[type, rank, DimList(1, 2, 3), DimList.create_unknown[rank](), alignment=4]

Args:

  • other (NDBuffer[type, rank, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): The other NDBuffer type.

__init__(ptr: UnsafePointer[scalar<#lit.struct.extract<:_stdlib::_builtin::_dtype::_DType type, "value">>, address_space=address_space], dynamic_shape: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> Self

Constructs an NDBuffer with statically known rank, but dynamic shapes and type.

Constraints:

The rank is known.

Args:

  • ptr (UnsafePointer[scalar<#lit.struct.extract<:_stdlib::_builtin::_dtype::_DType type, "value">>, address_space=address_space]): Pointer to the data.
  • dynamic_shape (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): A static tuple of size 'rank' representing shapes.

__init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space], dynamic_shape: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> Self

Constructs an NDBuffer with statically known rank, but dynamic shapes and type.

Constraints:

The rank is known.

Args:

  • ptr (UnsafePointer[SIMD[type, 1], address_space=address_space]): Pointer to the data.
  • dynamic_shape (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): A static tuple of size 'rank' representing shapes.

__init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space], dynamic_shape: DimList) -> Self

Constructs an NDBuffer with statically known rank, but dynamic shapes and type.

Constraints:

The rank is known.

Args:

  • ptr (UnsafePointer[SIMD[type, 1], address_space=address_space]): Pointer to the data.
  • dynamic_shape (DimList): A static tuple of size 'rank' representing shapes.

__init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space], dynamic_shape: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned], dynamic_stride: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> Self

Constructs a strided NDBuffer with statically known rank, but dynamic shapes and type.

Constraints:

The rank is known.

Args:

  • ptr (UnsafePointer[SIMD[type, 1], address_space=address_space]): Pointer to the data.
  • dynamic_shape (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): A static tuple of size 'rank' representing shapes.
  • dynamic_stride (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): A static tuple of size 'rank' representing strides.

__init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space], dynamic_shape: DimList, dynamic_stride: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> Self

Constructs a strided NDBuffer with statically known rank, but dynamic shapes and type.

Constraints:

The rank is known.

Args:

  • ptr (UnsafePointer[SIMD[type, 1], address_space=address_space]): Pointer to the data.
  • dynamic_shape (DimList): A DimList of size 'rank' representing shapes.
  • dynamic_stride (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): A static tuple of size 'rank' representing strides.

__getitem__

__getitem__(self, *idx: Int) -> SIMD[type, 1]

Gets an element from the buffer from the specified index.

Args:

  • *idx (Int): Index of the element to retrieve.

Returns:

The value of the element.

__getitem__(self, idx: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> SIMD[type, 1]

Gets an element from the buffer from the specified index.

Args:

  • idx (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): Index of the element to retrieve.

Returns:

The value of the element.

__setitem__

__setitem__(self, idx: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned], val: SIMD[type, 1])

Stores a single value into the buffer at the specified index.

Args:

  • idx (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): The index into the Buffer.
  • val (SIMD[type, 1]): The value to store.

__setitem__(self, *idx: Int, *, val: SIMD[type, 1])

Stores a single value into the buffer at the specified index.

Args:

  • *idx (Int): Index of the element to retrieve.
  • val (SIMD[type, 1]): The value to store.

get_rank

get_rank(self) -> Int

Returns the rank of the buffer.

Returns:

The rank of NDBuffer.

get_shape

get_shape(self) -> IndexList[rank]

Returns the shapes of the buffer.

Returns:

A static tuple of size 'rank' representing shapes of the NDBuffer.

get_strides

get_strides(self) -> IndexList[rank]

Returns the strides of the buffer.

Returns:

A static tuple of size 'rank' representing strides of the NDBuffer.

get_nd_index

get_nd_index(self, idx: Int) -> IndexList[rank]

Computes the NDBuffer's ND-index based on the flat index.

Args:

  • idx (Int): The flat index.

Returns:

The index positions.

__len__

__len__(self) -> Int

Computes the NDBuffer's number of elements.

Returns:

The total number of elements in the NDBuffer.

num_elements

num_elements(self) -> Int

Computes the NDBuffer's number of elements.

Returns:

The total number of elements in the NDBuffer.

size

size(self) -> Int

Computes the NDBuffer's number of elements.

Returns:

The total number of elements in the NDBuffer.

__str__

__str__(self) -> String

Gets the buffer as a string.

Returns:

A compact string of the buffer.

write_to

write_to[W: Writer](self, mut writer: W)

Formats this buffer to the provided Writer.

Parameters:

  • W (Writer): A type conforming to the Writable trait.

Args:

  • writer (W): The object to write to.

__repr__

__repr__(self) -> String

Gets the buffer as a string.

Returns:

A compact string representation of the buffer.

tile

tile[*tile_sizes: Dim](self, tile_coords: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> NDBuffer[type, rank, DimList(VariadicList(tile_sizes)), address_space=address_space]

Returns an n-d tile "slice" of the buffer of size tile_sizes at coords.

Parameters:

  • *tile_sizes (Dim): The size of the tiles.

Args:

  • tile_coords (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): The tile index.

Returns:

The tiled buffer at tile_coords.

load

load[*, width: Int = 1, alignment: Int = _default_alignment[::Int]()](self, *idx: Int) -> SIMD[type, width]

Loads a simd value from the buffer at the specified index.

Constraints:

The buffer must be contiguous or width must be 1.

Parameters:

  • width (Int): The simd_width of the load.
  • alignment (Int): The alignment value.

Args:

  • *idx (Int): The index into the NDBuffer.

Returns:

The simd value starting at the idx position and ending at idx+width.

load[*, width: Int = 1, alignment: Int = _default_alignment[::Int]()](self, idx: VariadicList[Int]) -> SIMD[type, width]

Loads a simd value from the buffer at the specified index.

Constraints:

The buffer must be contiguous or width must be 1.

Parameters:

  • width (Int): The simd_width of the load.
  • alignment (Int): The alignment value.

Args:

  • idx (VariadicList[Int]): The index into the NDBuffer.

Returns:

The simd value starting at the idx position and ending at idx+width.

load[*, width: Int = 1, alignment: Int = _default_alignment[::Int]()](self, idx: IndexList[size, element_bitwidth=element_bitwidth, unsigned=unsigned]) -> SIMD[type, width]

Loads a simd value from the buffer at the specified index.

Constraints:

The buffer must be contiguous or width must be 1.

Parameters:

  • width (Int): The simd_width of the load.
  • alignment (Int): The alignment value.

Args:

  • idx (IndexList[size, element_bitwidth=element_bitwidth, unsigned=unsigned]): The index into the NDBuffer.

Returns:

The simd value starting at the idx position and ending at idx+width.

load[*, width: Int = 1, alignment: Int = _default_alignment[::Int]()](self, idx: StaticTuple[Int, rank]) -> SIMD[type, width]

Loads a simd value from the buffer at the specified index.

Constraints:

The buffer must be contiguous or width must be 1.

Parameters:

  • width (Int): The simd_width of the load.
  • alignment (Int): The alignment value.

Args:

  • idx (StaticTuple[Int, rank]): The index into the NDBuffer.

Returns:

The simd value starting at the idx position and ending at idx+width.

store

store[*, width: Int = 1, alignment: Int = _default_alignment[::Int]()](self, idx: IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned], val: SIMD[type, width])

Stores a simd value into the buffer at the specified index.

Constraints:

The buffer must be contiguous or width must be 1.

Parameters:

  • width (Int): The width of the simd vector.
  • alignment (Int): The alignment value.

Args:

  • idx (IndexList[rank, element_bitwidth=element_bitwidth, unsigned=unsigned]): The index into the Buffer.
  • val (SIMD[type, width]): The value to store.

store[*, width: Int = 1, alignment: Int = _default_alignment[::Int]()](self, idx: StaticTuple[Int, rank], val: SIMD[type, width])

Stores a simd value into the buffer at the specified index.

Constraints:

The buffer must be contiguous or width must be 1.

Parameters:

  • width (Int): The width of the simd vector.
  • alignment (Int): The alignment value.

Args:

  • idx (StaticTuple[Int, rank]): The index into the Buffer.
  • val (SIMD[type, width]): The value to store.

dim

dim[index: Int](self) -> Int

Gets the buffer dimension at the given index.

Parameters:

  • index (Int): The number of dimension to get.

Returns:

The buffer size at the given dimension.

dim(self, index: Int) -> Int

Gets the buffer dimension at the given index.

Args:

  • index (Int): The number of dimension to get.

Returns:

The buffer size at the given dimension.

stride

stride[index: Int](self) -> Int

Gets the buffer stride at the given index.

Parameters:

  • index (Int): The number of dimension to get the stride for.

Returns:

The stride at the given dimension.

stride(self, index: Int) -> Int

Gets the buffer stride at the given index.

Args:

  • index (Int): The number of dimension to get the stride for.

Returns:

The stride at the given dimension.

is_contiguous

is_contiguous(self) -> Bool

Checks if the buffer is contiguous in memory.

Returns:

True if the buffer is contiguous in memory and False otherwise.

flatten

flatten(self) -> Buffer[type, shape.product(), address_space=address_space]

Constructs a flattened Buffer counterpart for this NDBuffer.

Constraints:

The buffer must be contiguous.

Returns:

Constructed Buffer object.

make_dims_unknown

make_dims_unknown(self) -> NDBuffer[type, rank, address_space=address_space]

Rebinds the NDBuffer to one with unknown shape.

Returns:

The rebound NDBuffer with unknown shape.

bytecount

bytecount(self) -> Int

Returns the size of the NDBuffer in bytes.

Returns:

The size of the NDBuffer in bytes.

zero

zero(self)

Sets all bytes of the NDBuffer to 0.

Constraints:

The buffer must be contiguous.

tofile

tofile(self, path: Path)

Write values to a file.

Args:

  • path (Path): Path to the output file.

fill

fill(self, val: SIMD[type, 1])

Assigns val to all elements in the Buffer.

The fill is performed in chunks of size N, where N is the native SIMD width of type on the system.

Args:

  • val (SIMD[type, 1]): The value to store.

stack_allocation

static stack_allocation[*, alignment: Int = alignof[::DType,__mlir_type.!kgen.target]().value]() -> Self

Constructs an NDBuffer instance backed by stack allocated memory space.

Parameters:

  • alignment (Int): Address alignment requirement for the allocation.

Returns:

Constructed NDBuffer with the allocated space.

prefetch

prefetch[params: PrefetchOptions](self, *idx: Int)

Prefetches the data at the given index.

Parameters:

  • params (PrefetchOptions): The prefetch configuration.

Args:

  • *idx (Int): The N-D index of the prefetched location.

prefetch[params: PrefetchOptions](self, indices: IndexList[rank])

Prefetches the data at the given index.

Parameters:

  • params (PrefetchOptions): The prefetch configuration.

Args:

  • indices (IndexList[rank]): The N-D index of the prefetched location.