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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!