Skip to main content

/

Docs

Nightly

v25.1

Log in

Mojo struct

Element

struct Element[dtype: DType, layout: Layout, /, *, bitwidth: Int = bitwidthof[AnyTrivialRegType,__mlir_type.!kgen.target]()]

A wrapper around SIMD types that provides layout-driven vectorized operations.

The Element struct extends SIMD types with layout-aware load and store operations, enabling efficient vectorized access to multi-dimensional data. It maps between logical tensor coordinates and physical memory locations according to the specified layout.

Parameters​

  • ​dtype (DType): The data type of the elements.
  • ​layout (Layout): The memory layout describing how elements are organized.
  • ​bitwidth (Int): The bit width used for runtime layout calculations.

Aliases​

  • element_data_type = SIMD[dtype, layout.size()]:

Fields​

  • ​element_data (SIMD[dtype, layout.size()]):
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=bitwidth]):

Implemented traits​

AnyType, Stringable, UnknownDestructibility, Writable

Methods​

__init__​

@implicit __init__(out self, element_data: SIMD[dtype, layout.size()])

Initializes an Element with the given SIMD data.

Args:

  • ​element_data (SIMD[dtype, layout.size()]): The SIMD data to initialize the element with.

__init__(out self, element_data: SIMD[dtype, layout.size()], runtime_layout: RuntimeLayout[layout, bitwidth=bitwidth])

Initializes an Element with the given SIMD data and runtime layout.

Args:

  • ​element_data (SIMD[dtype, layout.size()]): The SIMD data to initialize the element with.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=bitwidth]): The runtime layout to use for memory access.

load​

static load(ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], runtime_layout: RuntimeLayout[layout, bitwidth=bitwidth] = RuntimeLayout()) -> Self

Loads data from memory according to the specified layout.

This method loads data from memory using the layout information to determine the memory access pattern. It supports both rank-1 and rank-2 layouts with various stride patterns, optimizing for contiguous memory access when possible.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): Pointer to the memory location to load from.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=bitwidth]): The runtime layout to use for memory access.

Returns:

A new Element containing the loaded data.

masked_load​

static masked_load(ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], runtime_layout: RuntimeLayout[layout, bitwidth=bitwidth] = RuntimeLayout()) -> Self

Loads data from memory with masking for partial loads.

This method loads data from memory using the layout information, but also handles cases where the runtime dimensions are smaller than the static layout dimensions. It ensures that only valid memory locations are accessed.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): Pointer to the memory location to load from.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=bitwidth]): The runtime layout to use for memory access.

Returns:

A new Element containing the loaded data, with zeros in positions beyond the runtime dimensions.

store​

store(self, ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, origin=origin])

Stores element data to memory according to the specified layout.

This method performs a layout-aware store operation, writing data to memory following the access patterns defined by the layout. It optimizes memory writes based on the layout's stride patterns to maximize performance.

The method handles different memory layout patterns:

  • For rank-1 tensors with contiguous memory (stride=1), it uses vectorized stores
  • For rank-2 tensors with contiguous rows or columns, it uses optimized slice-based stores
  • For non-contiguous memory layouts, it performs element-by-element stores

Unlike masked_store(), this method assumes the full static dimensions will be written and does not perform runtime dimension boundary checking.

Note: This method is constrained to layouts with rank <= 2. For higher-rank tensors, consider decomposing the operation.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, origin=origin]): Mutable pointer to the memory location where data will be stored.

masked_store​

masked_store(self, ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, origin=origin])

Stores element data to memory with masking for partial stores.

This method performs a layout-aware store operation with boundary checking. It ensures that only valid memory locations are written to when the runtime dimensions are smaller than the static layout dimensions, preventing out-of-bounds memory access.

The method optimizes for different memory layouts:

  • For contiguous memory (stride=1), it uses vectorized stores when possible
  • For non-contiguous memory, it performs element-by-element stores
  • For all patterns, it respects runtime dimension bounds

Note: This method is constrained to layouts with rank <= 2. For higher-rank tensors, consider decomposing the operation.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, origin=origin]): Pointer to the memory location where data will be stored.

__str__​

__str__(self) -> String

Returns a string representation of the element.

Returns:

A string representation of the element's data.

write_to​

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

Writes the element to the specified writer.

Parameters:

  • ​W (Writer): Type parameter representing a Writer implementation.

Args:

  • ​writer (W): The writer to output the element representation to.