IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

Element

struct Element[dtype: DType, layout: Layout, /, index_type: DType = _get_index_type(layout)]

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.
  • ​index_type (DType): The integer type of the index pointing to each element.

Fields​

  • ​element_data (Element[dtype, layout, index_type].element_data_type): The actual SIMD data stored in this element. This field contains the vectorized data values that can be processed efficiently using SIMD operations.
  • ​runtime_layout (RuntimeLayout[layout, element_type=DType.int32, linear_idx_type=index_type]): The runtime layout information for memory access patterns. This field stores the layout information needed to map between logical tensor coordinates and physical memory locations, supporting both compile-time and runtime-determined access patterns.

Implemented traits​

AnyType, ImplicitlyDeletable, Writable

comptime members​

element_data_type​

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

The SIMD type used to store and process the element data.

This type alias defines a SIMD vector with the specified data type and size matching the layout's total element count, enabling efficient vectorized operations.

Methods​

__init__​

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

Initializes an Element with the given SIMD data.

Args:

def __init__(out self, element_data: SIMD[dtype, layout.size()], runtime_layout: RuntimeLayout[layout, element_type=DType.int32, linear_idx_type=index_type])

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

Args:

load​

static def load(ptr: Pointer[Scalar[dtype], address_space=ptr.address_space], runtime_layout: RuntimeLayout[layout, element_type=DType.int32, linear_idx_type=index_type] = 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:

Returns:

Self: A new Element containing the loaded data.

masked_load​

static def masked_load(ptr: Pointer[Scalar[dtype], address_space=ptr.address_space], runtime_layout: RuntimeLayout[layout, element_type=DType.int32, linear_idx_type=index_type] = 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:

Returns:

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

store​

def store(self, ptr: Pointer[Scalar[dtype], address_space=ptr.address_space])

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 (Pointer[Scalar[dtype], address_space=ptr.address_space]): Mutable pointer to the memory location where data will be stored.

masked_store​

def masked_store(self, ptr: Pointer[Scalar[dtype], address_space=ptr.address_space])

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 (Pointer[Scalar[dtype], address_space=ptr.address_space]): Pointer to the memory location where data will be stored.

write_to​

def write_to(self, mut writer: T)

Writes the element to the specified writer.

Args:

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

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact