Skip to main content

/

Docs

Log in

Mojo struct

Buffer

Defines a Buffer which can be parametrized on a static size and Dtype.

The Buffer does not own its underlying pointer.

Parameters​

  • ​type (DType): The element type of the Buffer.
  • ​size (Dim): The static size (if known) of the Buffer.
  • ​address_space (AddressSpace): The address space of the Buffer.

Fields​

  • ​data (UnsafePointer[SIMD[type, 1], address_space, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): The underlying data pointer of the data.
  • ​dynamic_size (Int): The dynamic size of the buffer.
  • ​dtype (DType): The dynamic data type of the buffer.

Implemented traits​

AnyType, Copyable, Movable, Sized

Methods​

__init__​

__init__(inout self: Self)

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

__init__(inout self: Self, ptr: UnsafePointer[SIMD[type, 1], address_space, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1])

Constructs a Buffer with statically known size and type.

Constraints:

The size is known.

Args:

  • ​ptr (UnsafePointer[SIMD[type, 1], address_space, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): Pointer to the data.

__init__(inout self: Self, ptr: UnsafePointer[SIMD[type, 1], address_space, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1], in_size: Int)

Constructs a Buffer with statically known type.

Constraints:

The size is unknown.

Args:

  • ​ptr (UnsafePointer[SIMD[type, 1], address_space, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): Pointer to the data.
  • ​in_size (Int): Dynamic size of the buffer.

__getitem__​

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

Loads a single element (SIMD of size 1) from the buffer at the specified index.

Args:

  • ​idx (Int): The index into the Buffer.

Returns:

The value at the idx position.

__setitem__​

__setitem__(self: Self, idx: Int, val: scalar<#lit.struct.extract<:_stdlib::_builtin::_dtype::_DType type, "value">>)

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

Args:

  • ​idx (Int): The index into the Buffer.
  • ​val (scalar<#lit.struct.extract<:_stdlib::_builtin::_dtype::_DType type, "value">>): The value to store.

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

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

Args:

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

__len__​

__len__(self: Self) -> Int

Gets the size if it is a known constant, otherwise it gets the dynamic_size.

This method is used by Buffer.__len__ to get the size of the buffer. If the Buffer size is a known constant, then the size is returned. Otherwise, the dynamic_size is returned.

Returns:

The size if static otherwise dynamic_size.

load​

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

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

Parameters:

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

Args:

  • ​idx (Int): The index into the Buffer.

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: Self, idx: Int, val: SIMD[type, width])

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

Parameters:

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

Args:

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

prefetch​

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

Prefetches the data at the given index.

Parameters:

  • ​params (PrefetchOptions): The prefetch configuration.

Args:

  • ​idx (Int): The index of the prefetched location.

bytecount​

bytecount(self: Self) -> Int

Returns the size of the Buffer in bytes.

Returns:

The size of the Buffer in bytes.

zero​

zero(self: Self)

Sets all bytes of the Buffer to 0.

fill​

fill(self: 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.

tofile​

tofile(self: Self, path: Path)

Write values to a file.

Args:

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

stack_allocation​

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

Constructs a buffer instance backed by stack allocated memory space.

Parameters:

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

Returns:

Constructed buffer with the allocated space.

Was this page helpful?