Skip to main content

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 (DTypePointer[type, address_space]): 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: LegacyPointer[SIMD[type, 1], address_space])

Constructs a Buffer with statically known size and type.

Constraints:

The size is known.

Args:

  • ptr (LegacyPointer[SIMD[type, 1], address_space]): Pointer to the data.

__init__(inout self: Self, /, ptr: DTypePointer[type, address_space])

Constructs a Buffer with statically known size and type.

Constraints:

The size is known.

Args:

  • ptr (DTypePointer[type, address_space]): Pointer to the data.

__init__(inout self: Self, /, ptr: LegacyPointer[SIMD[type, 1], address_space], in_size: Int)

Constructs a Buffer with statically known type.

Constraints:

The size is unknown.

Args:

  • ptr (LegacyPointer[SIMD[type, 1], address_space]): Pointer to the data.
  • in_size (Int): Dynamic size of the buffer.

__init__(inout self: Self, /, ptr: DTypePointer[type, address_space], in_size: Int)

Constructs a Buffer with statically known type.

Constraints:

The size is unknown.

Args:

  • ptr (DTypePointer[type, address_space]): 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 = alignof[stdlib::builtin::dtype::DType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1](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 = alignof[stdlib::builtin::dtype::DType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1](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.

simd_nt_store

simd_nt_store[width: Int](self: Self, idx: Int, val: SIMD[type, width])

Stores a simd value using non-temporal store.

Constraints:

The address must be properly aligned, 64B for avx512, 32B for avx2, and 16B for avx.

Parameters:

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

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.

aligned_stack_allocation

static aligned_stack_allocation[alignment: Int]() -> 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.

stack_allocation

static stack_allocation() -> Self

Constructs a buffer instance backed by stack allocated memory space.

Returns:

Constructed buffer with the allocated space.