Mojo struct
NDBuffer
@register_passable(trivial)
struct NDBuffer[mut: Bool, //, dtype: DType, rank: Int, origin: Origin[mut], shape: DimList = DimList.create_unknown[rank](), strides: DimList = DimList.create_unknown[rank](), *, alignment2: Int = 1, address_space: AddressSpace = AddressSpace.GENERIC, 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
- mut (
Bool): The inferred mutability. - dtype (
DType): The element dtype of the buffer. - rank (
Int): The rank of the buffer. - origin (
Origin): The origin of the memory being addressed. - shape (
DimList): The static size (if known) of the buffer. - strides (
DimList): The strides (if known) of the buffer. - alignment2 (
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 (
LegacyUnsafePointer[Scalar[dtype], address_space=address_space, mut=mut, origin=origin]): The underlying data for the buffer. The pointer is not owned by the NDBuffer. - dynamic_shape (
IndexList[rank, element_type=DType.uint64]): The dynamic value of the shape. - dynamic_stride (
IndexList[rank, element_type=DType.uint64]): The dynamic stride of the buffer.
Implemented traits
AnyType,
Copyable,
Defaultable,
DevicePassable,
ImplicitlyCopyable,
Movable,
Sized,
Stringable,
UnknownDestructibility,
Writable
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
__moveinit__is_trivial
alias __moveinit__is_trivial = True
device_type
alias device_type = NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]
OriginCastType
alias OriginCastType[target_mut: Bool, target_origin: Origin[target_mut]] = NDBuffer[dtype, rank, target_origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]
Parameters
type
alias type = dtype
Methods
__init__
__init__() -> Self
Default initializer for NDBuffer. By default the fields are all initialized to 0.
__init__(ptr: LegacyUnsafePointer[Scalar[dtype], address_space=address_space, mut=mut, origin=origin]) -> Self
Constructs an NDBuffer with statically known rank, shapes and dtype.
Constraints:
The rank, shapes, and type are known.
Args:
- ptr (
LegacyUnsafePointer): Pointer to the data.
__init__(span: Span[Scalar[dtype], origin, address_space=address_space]) -> Self
Constructs an NDBuffer with statically known rank, shapes and dtype.
Constraints:
The rank, shapes, and type are known.
Args:
- span (
Span): Span of the data.
@implicit
__init__(other: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, 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[dtype, rank, DimList(1, 2, 3), DimList(6, 6, 1), alignment=16]
to
NDBuffer[dtype, rank, DimList(1, 2, 3), DimList.create_unknown[rank](), alignment=4]
Args:
- other (
NDBuffer): The other NDBuffer type.
__init__(ptr: LegacyUnsafePointer[scalar<#lit.struct.extract<:_stdlib::_builtin::_dtype::_DType dtype, "_mlir_value">>, address_space=address_space, mut=mut, origin=origin], dynamic_shape: IndexList[rank, element_type=element_type]) -> Self
Constructs an NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
- ptr (
LegacyUnsafePointer): Pointer to the data. - dynamic_shape (
IndexList): A static tuple of size 'rank' representing shapes.
__init__(ptr: LegacyUnsafePointer[Scalar[dtype], address_space=address_space, mut=mut, origin=origin], dynamic_shape: IndexList[rank, element_type=element_type]) -> Self
Constructs an NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
- ptr (
LegacyUnsafePointer): Pointer to the data. - dynamic_shape (
IndexList): A static tuple of size 'rank' representing shapes.
__init__(span: Span[Scalar[dtype], origin, address_space=address_space], dynamic_shape: IndexList[rank, element_type=element_type]) -> Self
Constructs an NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
__init__(ptr: LegacyUnsafePointer[Scalar[dtype], address_space=address_space, mut=mut, origin=origin], dynamic_shape: DimList) -> Self
Constructs an NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
- ptr (
LegacyUnsafePointer): Pointer to the data. - dynamic_shape (
DimList): A static tuple of size 'rank' representing shapes.
__init__(span: Span[Scalar[dtype], origin, 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:
__init__(ptr: LegacyUnsafePointer[Scalar[dtype], address_space=address_space, mut=mut, origin=origin], dynamic_shape: IndexList[rank, element_type=element_type], dynamic_stride: IndexList[rank, element_type=element_type]) -> Self
Constructs a strided NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
- ptr (
LegacyUnsafePointer): Pointer to the data. - dynamic_shape (
IndexList): A static tuple of size 'rank' representing shapes. - dynamic_stride (
IndexList): A static tuple of size 'rank' representing strides.
__init__(span: Span[Scalar[dtype], origin, address_space=address_space], dynamic_shape: IndexList[rank, element_type=element_type], dynamic_stride: IndexList[rank, element_type=element_type]) -> Self
Constructs a strided NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
__init__(ptr: LegacyUnsafePointer[Scalar[dtype], address_space=address_space, mut=mut, origin=origin], dynamic_shape: DimList, dynamic_stride: IndexList[rank, element_type=element_type]) -> Self
Constructs a strided NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
- ptr (
LegacyUnsafePointer): Pointer to the data. - dynamic_shape (
DimList): A DimList of size 'rank' representing shapes. - dynamic_stride (
IndexList): A static tuple of size 'rank' representing strides.
__init__(span: Span[Scalar[dtype], origin, address_space=address_space], dynamic_shape: DimList, dynamic_stride: IndexList[rank, element_type=element_type]) -> Self
Constructs a strided NDBuffer with statically known rank, but dynamic shapes and type.
Constraints:
The rank is known.
Args:
__getitem__
__getitem__(self, *idx: Int) -> Scalar[dtype]
Gets an element from the buffer from the specified index.
Args:
- *idx (
Int): Index of the element to retrieve.
Returns:
Scalar: The value of the element.
__getitem__(self, idx: IndexList[rank, element_type=element_type]) -> Scalar[dtype]
Gets an element from the buffer from the specified index.
Args:
- idx (
IndexList): Index of the element to retrieve.
Returns:
Scalar: The value of the element.
__setitem__
__setitem__(self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive], idx: IndexList[rank, element_type=element_type], val: Scalar[dtype])
Stores a single value into the buffer at the specified index.
Args:
__setitem__(self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive], *idx: Int, *, val: Scalar[dtype])
Stores a single value into the buffer at the specified index.
Args:
get_immutable
get_immutable(self) -> NDBuffer[dtype, rank, origin_of((muttoimm origin._mlir_origin)), shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]
Changes the mutability of the NDBuffer to immutable.
Returns:
NDBuffer: A buffer with the mutability set to immutable.
as_any_origin
as_any_origin(self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]) -> NDBuffer[dtype, rank, MutAnyOrigin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]
Changes the origin of the NDBuffer to MutAnyOrigin.
This requires the buffer to already be mutable as casting mutability is inherently very unsafe.
It is usually preferred to maintain concrete origin values instead of
using MutAnyOrigin. However, if it is needed, keep in mind that
MutAnyOrigin can alias any memory value, so Mojo's ASAP
destruction will not apply during the lifetime of the buffer.
Returns:
NDBuffer: A buffer with the origin set to MutAnyOrigin.
as_any_origin(self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]) -> NDBuffer[dtype, rank, ImmutAnyOrigin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive]
Changes the origin of the NDBuffer to ImmutAnyOrigin.
It is usually preferred to maintain concrete origin values instead of
using ImmutAnyOrigin. However, if it is needed, keep in mind that
ImmutAnyOrigin can alias any memory value, so Mojo's ASAP
destruction will not apply during the lifetime of the buffer.
Returns:
NDBuffer: A buffer with the origin set to ImmutAnyOrigin.
get_rank
get_shape
get_shape(self) -> IndexList[rank]
Returns the shapes of the buffer.
Returns:
IndexList: 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:
IndexList: 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:
IndexList: The index positions.
__len__
__len__(self) -> Int
Computes the NDBuffer's number of elements.
Returns:
Int: The total number of elements in the NDBuffer.
num_elements
num_elements(self) -> Int
Computes the NDBuffer's number of elements.
Returns:
Int: The total number of elements in the NDBuffer.
size
size(self) -> Int
Computes the NDBuffer's number of elements.
Returns:
Int: The total number of elements in the NDBuffer.
__str__
__str__(self) -> String
Gets the buffer as a string.
Returns:
String: A compact string of the buffer.
write_to
write_to(self, mut writer: T)
Formats this buffer to the provided Writer.
Args:
- writer (
T): The object to write to.
__repr__
__repr__(self) -> String
Gets the buffer as a string.
Returns:
String: A compact string representation of the buffer.
tile
tile[*tile_sizes: Dim](self, tile_coords: IndexList[rank, element_type=element_type]) -> NDBuffer[dtype, rank, origin, DimList(VariadicList[Dim](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): The tile index.
Returns:
NDBuffer: The tiled buffer at tile_coords.
load
load[*, width: Int = 1, alignment: Int = NDBuffer._default_alignment[mut, dtype, rank, origin, shape, strides, alignment2, address_space, exclusive, width]()](self, *idx: Int) -> SIMD[dtype, width]
Loads a simd value from the buffer at the specified index.
Constraints:
The buffer must be contiguous or width must be 1.
Parameters:
Args:
- *idx (
Int): The index into the NDBuffer.
Returns:
SIMD: The simd value starting at the idx position and ending at
idx+width.
load[*, width: Int = 1, alignment: Int = NDBuffer._default_alignment[mut, dtype, rank, origin, shape, strides, alignment2, address_space, exclusive, width]()](self, idx: VariadicList[Int]) -> SIMD[dtype, width]
Loads a simd value from the buffer at the specified index.
Constraints:
The buffer must be contiguous or width must be 1.
Parameters:
Args:
- idx (
VariadicList): The index into the NDBuffer.
Returns:
SIMD: The simd value starting at the idx position and ending at
idx+width.
load[*, width: Int = 1, alignment: Int = NDBuffer._default_alignment[mut, dtype, rank, origin, shape, strides, alignment2, address_space, exclusive, width]()](self, idx: IndexList[size, element_type=element_type]) -> SIMD[dtype, width]
Loads a simd value from the buffer at the specified index.
Constraints:
The buffer must be contiguous or width must be 1.
Parameters:
Args:
- idx (
IndexList): The index into the NDBuffer.
Returns:
SIMD: The simd value starting at the idx position and ending at
idx+width.
load[*, width: Int = 1, alignment: Int = NDBuffer._default_alignment[mut, dtype, rank, origin, shape, strides, alignment2, address_space, exclusive, width]()](self, idx: StaticTuple[Int, rank]) -> SIMD[dtype, width]
Loads a simd value from the buffer at the specified index.
Constraints:
The buffer must be contiguous or width must be 1.
Parameters:
Args:
- idx (
StaticTuple): The index into the NDBuffer.
Returns:
SIMD: The simd value starting at the idx position and ending at
idx+width.
store
store[_alignment: Int, //, *, width: Int = 1, alignment: Int = NDBuffer._default_alignment[mut, dtype, rank, origin, shape, strides, alignment2, address_space, exclusive, width]()](self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=_alignment, address_space=address_space, exclusive=exclusive], idx: IndexList[rank, element_type=element_type], val: SIMD[dtype, width])
Stores a simd value into the buffer at the specified index.
Constraints:
The buffer must be contiguous or width must be 1.
Parameters:
- _alignment (
Int): The inferred alignment of self. - width (
Int): The width of the simd vector. - alignment (
Int): The alignment value.
Args:
store[_alignment: Int, //, *, width: Int = 1, alignment: Int = NDBuffer._default_alignment[mut, dtype, rank, origin, shape, strides, alignment2, address_space, exclusive, width]()](self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=_alignment, address_space=address_space, exclusive=exclusive], idx: StaticTuple[Int, rank], val: SIMD[dtype, width])
Stores a simd value into the buffer at the specified index.
Constraints:
The buffer must be contiguous or width must be 1.
Parameters:
- _alignment (
Int): The inferred alignment of self. - width (
Int): The width of the simd vector. - alignment (
Int): The alignment value.
Args:
- idx (
StaticTuple): The index into the buffer. - val (
SIMD): 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:
Int: 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:
Int: 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:
Int: 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:
Int: The stride at the given dimension.
is_contiguous
is_contiguous(self) -> Bool
Checks if the buffer is contiguous in memory.
Returns:
Bool: True if the buffer is contiguous in memory and False otherwise.
flatten
flatten(self) -> NDBuffer[dtype, 1, origin, shape.product(), address_space=address_space]
Constructs a flattened buffer counterpart for this NDBuffer.
Constraints:
The buffer must be contiguous.
Returns:
NDBuffer: Constructed buffer object.
make_dims_unknown
make_dims_unknown(self) -> NDBuffer[dtype, rank, origin, address_space=address_space]
Rebinds the NDBuffer to one with unknown shape.
Returns:
NDBuffer: The rebound NDBuffer with unknown shape.
bytecount
bytecount(self) -> Int
Returns the size of the NDBuffer in bytes.
Returns:
Int: The size of the NDBuffer in bytes.
zero
zero(self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive])
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.
Raises:
If the operation fails.
fill
fill(self: NDBuffer[dtype, rank, origin, shape, strides, alignment2=alignment2, address_space=address_space, exclusive=exclusive], val: Scalar[dtype])
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 (
Scalar): The value to store.
stack_allocation
static stack_allocation[*, alignment: Int = align_of[dtype]()]() -> Self
Constructs an NDBuffer instance backed by stack allocated memory space.
Parameters:
- alignment (
Int): Address alignment requirement for the allocation.
Returns:
Self: 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): The N-D index of the prefetched location.
get_type_name
static get_type_name() -> String
Gets the name of the host type (the one implementing this trait).
Returns:
String: The host type's name.
get_device_type_name
static get_device_type_name() -> String
Gets device_type's name.
Returns:
String: The device type's name.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!