Mojo struct
LayoutTensorIter
@register_passable(trivial)
struct LayoutTensorIter[mut: Bool, //, type: DType, layout: Layout, origin: Origin[mut], /, *, address_space: AddressSpace = AddressSpace(0), alignment: Int = alignof[::DType,__mlir_type.!kgen.target]() if is_nvidia_gpu() else 1, circular: Bool = False, axis: OptionalReg[Int] = OptionalReg[Int]({:i1 0, 1}), layout_bitwidth: Int = bitwidthof[::DType,__mlir_type.!kgen.target](), 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.
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. - type (
DType): The data type of the tensor elements. - layout (
Layout): The memory layout pattern to follow during iteration. - origin (
Origin[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 (
OptionalReg[Int]): Optional axis for dimension-specific operations. - layout_bitwidth (
Int): Bit width used for layout indices. - masked (
Bool): Whether to apply bounds masking during iteration.
Aliases
uint_type = SIMD[_get_unsigned_type(layout, address_space), 1]:
Fields
- ptr (
UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): - offset (
SIMD[_get_unsigned_type(layout, address_space), 1]): - stride (
SIMD[_get_unsigned_type(layout, address_space), 1]): - bound (
SIMD[_get_unsigned_type(layout, address_space), 1]): - runtime_layout (
RuntimeLayout[layout, bitwidth=layout_bitwidth]): - dimension_bound (
SIMD[_get_unsigned_type(layout, address_space), 1]): - idx (
SIMD[_get_unsigned_type(layout, address_space), 1]):
Implemented traits
AnyType,
UnknownDestructibility
Methods
__init__
__init__() -> Self
Initialize an empty iterator.
Creates a default iterator with zero values, typically used as a placeholder or default value.
__init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], bound: SIMD[_get_unsigned_type(layout, address_space), 1], stride: SIMD[_get_unsigned_type(layout, address_space), 1] = SIMD(layout.size()), offset: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!pop.int_literal](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:
- ptr (
UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): Pointer to the beginning of the memory region. - bound (
SIMD[_get_unsigned_type(layout, address_space), 1]): Upper bound of the memory region. - stride (
SIMD[_get_unsigned_type(layout, address_space), 1]): Step size between consecutive elements (defaults to layout size). - offset (
SIMD[_get_unsigned_type(layout, address_space), 1]): Initial offset from the base pointer.
__init__(ptr: UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], bound: SIMD[_get_unsigned_type(layout, address_space), 1], runtime_layout: RuntimeLayout[layout, bitwidth=bitwidth], stride: SIMD[_get_unsigned_type(layout, address_space), 1] = SIMD(layout.size() if layout.all_dims_known() else -1), offset: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!pop.int_literal](0), dimension_bound: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!pop.int_literal](0), idx: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!pop.int_literal](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:
- ptr (
UnsafePointer[SIMD[type, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): Pointer to the beginning of the memory region. - bound (
SIMD[_get_unsigned_type(layout, address_space), 1]): Upper bound of the memory region. - runtime_layout (
RuntimeLayout[layout, bitwidth=bitwidth]): Layout determined at runtime. - stride (
SIMD[_get_unsigned_type(layout, address_space), 1]): Step size between consecutive elements. - offset (
SIMD[_get_unsigned_type(layout, address_space), 1]): Initial offset from the base pointer. - dimension_bound (
SIMD[_get_unsigned_type(layout, address_space), 1]): Bound for the specified dimension when using masked iteration. - idx (
SIMD[_get_unsigned_type(layout, address_space), 1]): Initial index position.
__getitem__
__getitem__(self) -> LayoutTensor[type, layout, origin, address_space=address_space, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]
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:
A layout tensor at the current iterator position.
__iadd__
__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.
This function is unsafe. It omits bound checking for performance reasons.
Caller must ensure the index doesn't go out-of-bounds.
Args:
- rhs (
T): The number of positions to advance.
__iadd__(mut self, rhs: SIMD[_get_unsigned_type(layout, address_space), 1])
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.
This function is unsafe. It omits bound checking for performance reasons.
Caller must ensure the index doesn't go out-of-bounds.
Args:
- rhs (
SIMD[_get_unsigned_type(layout, address_space), 1]): The number of positions to advance.
get
get(self) -> LayoutTensor[type, layout, origin, address_space=address_space, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]
Get the layout tensor at the current iterator position.
Returns a layout tensor representing the data at the current position of the iterator.
next
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.
Args:
- rhs (
T): The number of positions to advance.
Returns:
A new iterator pointing to the advanced position.
next(self, rhs: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!pop.int_literal](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:
- rhs (
SIMD[_get_unsigned_type(layout, address_space), 1]): The number of positions to advance (defaults to 1).
Returns:
A new iterator pointing to the advanced position.
next_unsafe
next_unsafe(self, rhs: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!pop.int_literal](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:
- rhs (
SIMD[_get_unsigned_type(layout, address_space), 1]): The number of positions to advance (defaults to 1).
Returns:
A new iterator pointing to the advanced position.
reshape
reshape[dst_layout: Layout](self) -> LayoutTensorIter[type, dst_layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_bitwidth=layout_bitwidth, 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:
A new iterator with the specified layout.
bitcast
bitcast[new_type: DType, *, address_space: AddressSpace = address_space, alignment: Int = alignment](self) -> LayoutTensorIter[new_type, layout, origin, address_space=address_space, alignment=alignment, circular=circular, layout_bitwidth=layout_bitwidth, 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. - address_space (
AddressSpace): The memory address space for the new iterator (defaults to current). - alignment (
Int): Memory alignment requirement for the new iterator (defaults to current).
Returns:
A new LayoutTensorIter with the same layout but different data type.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!