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

LayoutTensorIter

struct LayoutTensorIter[mut: Bool, //, dtype: DType, layout: Layout, origin: Origin[mut=mut], /, *, address_space: AddressSpace = AddressSpace.GENERIC, alignment: Int = align_of[dtype](), circular: Bool = False, axis: Optional[Int] = None, layout_int_type: DType = _get_index_type(address_space), linear_idx_type: DType = _get_index_type(address_space), masked: Bool = False]

Iterator for traversing a memory buffer with a specific layout.

LayoutTensorIter provides a way to iterate through memory according to a specific layout pattern, constructing layout tensors at each position. This enables efficient traversal of multi-dimensional data structures with custom memory layouts.

Notes:

The returned layout tensor is NOT vectorized. Users should explicitly vectorize if needed for performance-critical operations.

Parameters​

  • ​mut (Bool): Whether the iterator allows mutation of the underlying data.
  • ​dtype (DType): The data type of the tensor elements.
  • ​layout (Layout): The memory layout pattern to follow during iteration.
  • ​origin (Origin[mut=mut]): Origin tracking for memory safety.
  • ​address_space (AddressSpace): The memory address space (GLOBAL, SHARED, etc.).
  • ​alignment (Int): Memory alignment requirement for the data.
  • ​circular (Bool): Whether iteration wraps around at boundaries.
  • ​axis (Optional[Int]): Optional axis for dimension-specific operations.
  • ​layout_int_type (DType): Integer type used for layout indices.
  • ​linear_idx_type (DType): Integer type used for indexing into memory.
  • ​masked (Bool): Whether to apply bounds masking during iteration.

Fields​

  • ​ptr (Pointer[Scalar[dtype], origin, address_space=address_space]): Pointer to the memory region being iterated, with appropriate type and memory attributes.
  • ​offset (LayoutTensorIter[dtype, layout, origin, address_space=address_space, alignment=alignment, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked].linear_uint_type): Current offset from the base pointer, representing the iterator's position in memory.
  • ​stride (LayoutTensorIter[dtype, layout, origin, address_space=address_space, alignment=alignment, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked].linear_uint_type): Step size between consecutive elements or blocks in memory during iteration.
  • ​bound (LayoutTensorIter[dtype, layout, origin, address_space=address_space, alignment=alignment, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked].linear_uint_type): Upper bound of the memory region, limiting the iteration range.
  • ​runtime_layout (LayoutTensorIter[dtype, layout, origin, address_space=address_space, alignment=alignment, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked].RuntimeLayoutType): Runtime representation of the layout pattern used for mapping logical indices to memory locations.
  • ​dimension_bound (LayoutTensorIter[dtype, layout, origin, address_space=address_space, alignment=alignment, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked].layout_uint_type): Boundary value for the current dimension when iterating along a specific axis.
  • ​idx (LayoutTensorIter[dtype, layout, origin, address_space=address_space, alignment=alignment, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked].linear_uint_type): Current logical index position within the iteration sequence.

Implemented traits​

AnyType, Copyable, Defaultable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

BitcastType​

comptime BitcastType[new_type: DType, *, address_space: AddressSpace = address_space, alignment: Int = alignment] = LayoutTensorIter[new_type, layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked]

Type alias for bitcast iterator types.

Parameters​

  • ​new_type (DType): The target data type.
  • ​address_space (AddressSpace): The target address space.
  • ​alignment (Int): The target memory alignment.

layout_uint_type​

comptime layout_uint_type = Scalar[_unsigned_integral_type_of[layout_int_type]()]

The unsigned integer type used for layout, based on layout and address space.

LayoutTensorType​

comptime LayoutTensorType = LayoutTensor[dtype, layout, origin, address_space=address_space, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]

The LayoutTensor type returned by this iterator.

linear_uint_type​

comptime linear_uint_type = Scalar[_unsigned_integral_type_of[linear_idx_type]()]

The unsigned integer type used for indexing into memory.

ReshapeType​

comptime ReshapeType[dst_layout: Layout] = LayoutTensorIter[dtype, dst_layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked]

Type alias for reshaped iterator types.

Parameters​

  • ​dst_layout (Layout): The target layout for the reshaped iterator.

RuntimeLayoutType​

comptime RuntimeLayoutType = RuntimeLayout[layout, element_type=layout_int_type, linear_idx_type=linear_idx_type]

Type alias for the runtime layout.

Methods​

__init__​

def __init__() -> Self

Initialize an empty iterator.

Creates a default iterator with zero values, typically used as a placeholder or default value.

def __init__(ptr: Pointer[Scalar[dtype], origin, address_space=address_space], bound: Scalar[_unsigned_integral_type_of[linear_idx_type]()], stride: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = SIMD(layout.size()), offset: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = 0) -> Self

Initialize an iterator with a pointer and basic parameters.

Creates an iterator for a memory region with the specified bounds and stride.

Constraints:

The layout must have all dimensions known at compile time.

Args:

def __init__(ptr: Pointer[Scalar[dtype], origin, address_space=address_space], bound: Int) -> Self

Initialize an iterator with a pointer and Int bound.

Creates an iterator for a memory region with the specified bounds and stride.

Args:

  • ​ptr (Pointer[Scalar[dtype], origin, address_space=address_space]): Pointer to the beginning of the memory region.
  • ​bound (Int): Upper bound of the memory region.

def __init__(ptr: Pointer[Scalar[dtype], origin, address_space=address_space], bound: Scalar[_unsigned_integral_type_of[linear_idx_type]()], runtime_layout: RuntimeLayout[layout, element_type=runtime_layout.element_type, linear_idx_type=runtime_layout.linear_idx_type], stride: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = SIMD(layout.size() if layout.all_dims_known() else Int(-1)), offset: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = 0, dimension_bound: Scalar[_unsigned_integral_type_of[layout_int_type]()] = 0, idx: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = 0) -> Self

Initialize an iterator with a runtime layout.

Creates an iterator with a runtime-determined layout, allowing for more flexible memory traversal patterns.

Constraints:

The runtime layout must have the same bitwidth as specified for the iterator. Circular iteration is not supported when an axis is defined.

Args:

@implicit def __init__(other: LayoutTensorIter[address_space=other.address_space, alignment=other.alignment, circular=other.circular, axis=other.axis, layout_int_type=other.layout_int_type, linear_idx_type=other.linear_idx_type, masked=other.masked]) -> LayoutTensorIter[other.dtype, other.layout, origin_of(other.origin), address_space=other.address_space, alignment=other.alignment, circular=other.circular, axis=other.axis, layout_int_type=other.layout_int_type, linear_idx_type=other.linear_idx_type, masked=other.masked]

Implicitly cast a mutable LayoutTensorIter to immutable.

Args:

Returns:

LayoutTensorIter[other.dtype, other.layout, origin_of(other.origin), address_space=other.address_space, alignment=other.alignment, circular=other.circular, axis=other.axis, layout_int_type=other.layout_int_type, linear_idx_type=other.linear_idx_type, masked=other.masked]

__getitem__​

def __getitem__(self) -> Self.LayoutTensorType

Get the layout tensor at the current iterator position.

Operator overload that returns a layout tensor representing the data at the current position of the iterator.

Returns:

Self.LayoutTensorType: A layout tensor at the current iterator position.

__iadd__​

def __iadd__[T: Intable](mut self, rhs: T)

Increment the iterator by an integer value.

Advances the iterator by the specified number of positions.

Notes:

This function is unsafe. It omits bound checking for performance reasons. Caller must ensure the index doesn't go out-of-bounds.

Parameters:

  • ​T (Intable): A type that can be converted to an integer.

Args:

  • ​rhs (T): The number of positions to advance.

def __iadd__(mut self, rhs: Scalar[_unsigned_integral_type_of[linear_idx_type]()])

Increment the iterator by a uint value.

Advances the iterator by the specified number of positions.

Notes:

This function is unsafe. It omits bound checking for performance reasons. Caller must ensure the index doesn't go out-of-bounds.

Args:

get​

def get(self) -> Self.LayoutTensorType

Get the layout tensor at the current iterator position.

Returns a layout tensor representing the data at the current position of the iterator.

Returns:

Self.LayoutTensorType: A tensor view at the current iterator position with the same type, layout, and memory characteristics as specified by the output parameter.

next​

def next[T: Intable](self, rhs: T) -> Self

Return an iterator pointing to a position ahead by rhs steps.

Creates a new iterator that points rhs positions ahead of the current one.

Parameters:

  • ​T (Intable): An integer-convertible type for the step size.

Args:

  • ​rhs (T): The number of positions to advance.

Returns:

Self: A new iterator pointing to the advanced position.

def next(self, rhs: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = 1) -> Self

Return an iterator pointing to a position ahead by rhs steps.

Creates a new iterator that points rhs positions ahead of the current one.

Args:

Returns:

Self: A new iterator pointing to the advanced position.

next_unsafe​

def next_unsafe(self, rhs: Scalar[_unsigned_integral_type_of[linear_idx_type]()] = 1) -> Self

Return an iterator pointing to a position ahead by rhs steps (unsafe version).

Creates a new iterator that points rhs positions ahead of the current one. This is an unsafe version that omits certain checks for performance.

Constraints:

Cannot be used with masked iterators. User must ensure rhs < bound / stride.

Args:

Returns:

Self: A new iterator pointing to the advanced position.

reshape​

def reshape[dst_layout: Layout](self) -> LayoutTensorIter[dtype, dst_layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked]

Reshape the iterator to a new layout.

This method creates a new iterator with a different layout while preserving the underlying data. The new layout must have the same total size as the original.

Constraints:

  • The destination layout must have the same total size as the original.
  • Both layouts must be contiguous.
  • Both layouts must have compile-time known dimensions.

Parameters:

  • ​dst_layout (Layout): The target layout to reshape to.

Returns:

LayoutTensorIter[dtype, dst_layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked]: A new iterator with the specified layout.

bitcast​

def bitcast[new_type: DType, *, target_address_space: AddressSpace = address_space, target_alignment: Int = alignment](self) -> LayoutTensorIter[new_type, layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked]

Reinterpret the iterator's underlying pointer as a different data type.

This method performs a bitcast operation, allowing you to view the same memory location as a different data type without copying or converting the data.

Parameters:

  • ​new_type (DType): The target data type to cast to.
  • ​target_address_space (AddressSpace): The memory address space for the new iterator (defaults to current).
  • ​target_alignment (Int): Memory alignment requirement for the new iterator (defaults to current).

Returns:

LayoutTensorIter[new_type, layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked]: A new LayoutTensorIter with the same layout but different data type.

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact