Mojo struct
LayoutTensorBuild
@register_passable(trivial)
struct LayoutTensorBuild[dtype: DType, *, __layout: Layout = __init__[::Origin[::Bool(IntTuple(1)), __layout_init: Bool = False, __address_space: AddressSpace = AddressSpace(0), __layout_bitwidth: Int = bitwidthof[::DType,__mlir_type.!kgen.target](), __index_type: DType = _get_index_type(__layout, __address_space), __circular: Bool = False]
Tensor layout builder providing a fluent interface for constructing tensors with various layouts.
Parameters
- dtype (
DType): Data type of tensor elements. - __layout (
Layout): The tensor's memory layout. - __layout_init (
Bool): Whether the layout has been initialized. - __address_space (
AddressSpace): Memory space (generic, shared, local). - __layout_bitwidth (
Int): Bit width for layout index type. - __index_type (
DType): Type used for indexing. - __circular (
Bool): Whether tensor has circular indexing semantics.
Fields
- runtime_layout (
RuntimeLayout[__layout, bitwidth=__layout_bitwidth, linear_idx_type=__index_type]): Runtime representation of the tensor's layout.
Implemented traits
AnyType,
Copyable,
ExplicitlyCopyable,
Movable,
UnknownDestructibility
Methods
__init__
__init__() -> Self
Initializes a new LayoutTensorBuild instance with default values.
row_major
row_major[*shapes: Int](self) -> LayoutTensorBuild[dtype, __layout=row_major[::Origin[::Bool(_to_int_tuple[::VariadicList[::Int]]()), __layout_init=True]
Creates a row-major layout using compile-time dimensions.
Parameters:
- *shapes (
Int): Variadic parameter specifying the dimensions of the tensor. Each value represents the size of a dimension.
Returns:
LayoutTensorBuild - A new builder with row-major layout.
row_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=row_major(dim, dim), __layout_init=True]
Creates a row-major 2D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size.
Returns:
LayoutTensorBuild - A new builder with row-major layout.
row_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim], shape2: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=row_major(dim, dim, dim), __layout_init=True]
Creates a row-major 3D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size. - shape2 (
ValueOrUnknown[dim]): Third dimension size.
Returns:
LayoutTensorBuild - A new builder with row-major layout.
row_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim], shape2: ValueOrUnknown[dim], shape3: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=row_major(dim, dim, dim, dim), __layout_init=True]
Creates a row-major 4D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size. - shape2 (
ValueOrUnknown[dim]): Third dimension size. - shape3 (
ValueOrUnknown[dim]): Fourth dimension size.
Returns:
LayoutTensorBuild - A new builder with row-major layout.
row_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim], shape2: ValueOrUnknown[dim], shape3: ValueOrUnknown[dim], shape4: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=row_major(dim, dim, dim, dim, dim), __layout_init=True]
Creates a row-major 5D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size. - shape2 (
ValueOrUnknown[dim]): Third dimension size. - shape3 (
ValueOrUnknown[dim]): Fourth dimension size. - shape4 (
ValueOrUnknown[dim]): Fifth dimension size.
Returns:
LayoutTensorBuild - A new builder with row-major layout.
col_major
col_major[*shapes: Int](self) -> LayoutTensorBuild[dtype, __layout=col_major[::Origin[::Bool(_to_int_tuple[::VariadicList[::Int]]()), __layout_init=True]
Creates a column-major layout using compile-time dimensions.
Parameters:
- *shapes (
Int): Variadic parameter specifying the dimensions of the tensor. Each value represents the size of a dimension.
Returns:
LayoutTensorBuild - A new builder with column-major layout.
col_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=col_major(dim, dim), __layout_init=True]
Creates a column-major 2D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size.
Returns:
LayoutTensorBuild - A new builder with column-major layout.
col_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim], shape2: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=col_major(dim, dim, dim), __layout_init=True]
Creates a column-major 3D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size. - shape2 (
ValueOrUnknown[dim]): Third dimension size.
Returns:
LayoutTensorBuild - A new builder with column-major layout.
col_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim], shape2: ValueOrUnknown[dim], shape3: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=col_major(dim, dim, dim, dim), __layout_init=True]
Creates a column-major 4D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size. - shape2 (
ValueOrUnknown[dim]): Third dimension size. - shape3 (
ValueOrUnknown[dim]): Fourth dimension size.
Returns:
LayoutTensorBuild - A new builder with column-major layout.
col_major(self, shape0: ValueOrUnknown[dim], shape1: ValueOrUnknown[dim], shape2: ValueOrUnknown[dim], shape3: ValueOrUnknown[dim], shape4: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=col_major(dim, dim, dim, dim, dim), __layout_init=True]
Creates a column-major 5D layout using runtime dimensions.
Args:
- shape0 (
ValueOrUnknown[dim]): First dimension size. - shape1 (
ValueOrUnknown[dim]): Second dimension size. - shape2 (
ValueOrUnknown[dim]): Third dimension size. - shape3 (
ValueOrUnknown[dim]): Fourth dimension size. - shape4 (
ValueOrUnknown[dim]): Fifth dimension size.
Returns:
LayoutTensorBuild - A new builder with column-major layout.
layout
layout[shape0: Int](self) -> LayoutTensorBuild[dtype, __layout=__init__[::Origin[::Bool(IntTuple(shape0)), __layout_init=True]
Creates a 1D layout with a compile-time dimension.
Parameters:
- shape0 (
Int): Size of the single dimension.
Returns:
LayoutTensorBuild - A new builder with the specified layout.
layout[rank: Int, shape: Index[rank], stride: Index[rank]](self) -> LayoutTensorBuild[dtype, __layout=__init__[::Origin[::Bool(_to_int_tuple[::Int](shape), _to_int_tuple[::Int](stride)), __layout_init=True]
Creates a custom layout with compile-time dimensions and strides.
Parameters:
- rank (
Int): Number of dimensions. - shape (
Index[rank]): List of dimension sizes. - stride (
Index[rank]): List of strides for each dimension.
Returns:
LayoutTensorBuild - A new builder with the specified custom layout.
layout[rank: Int](self, shape: Index[rank], stride: Index[rank]) -> LayoutTensorBuild[dtype, __layout=__init__[::Origin[::Bool(_to_int_tuple[::Int](-1), _to_int_tuple[::Int](-1)), __layout_init=True]
Creates a custom layout with runtime dimensions and strides.
Parameters:
- rank (
Int): Number of dimensions.
Args:
- shape (
Index[rank]): List of dimension sizes. - stride (
Index[rank]): List of strides for each dimension.
Returns:
LayoutTensorBuild - A new builder with the specified custom layout.
layout(self, shape0: ValueOrUnknown[dim]) -> LayoutTensorBuild[dtype, __layout=__init__[::Origin[::Bool(IntTuple(dim)), __layout_init=True]
Creates a 1D layout with a runtime dimension.
Args:
- shape0 (
ValueOrUnknown[dim]): Size of the single dimension.
Returns:
LayoutTensorBuild - A new builder with the specified layout.
shared
shared(self) -> LayoutTensorBuild[dtype, __layout=__layout, __layout_init=__layout_init, __address_space=AddressSpace(3)]
Places the tensor in GPU shared memory.
Returns:
LayoutTensorBuild - A new builder with shared memory address space.
local
local(self) -> LayoutTensorBuild[dtype, __layout=__layout, __layout_init=__layout_init, __address_space=AddressSpace(5)]
Places the tensor in GPU local memory.
Returns:
LayoutTensorBuild - A new builder with local memory address space.
alloc
alloc(self) -> LayoutTensor[dtype, __layout, MutableAnyOrigin, address_space=__address_space]
Allocates a new tensor using the current layout.
Note: Fails to compile if layout is not set, dimensions are not known, or tensor is circular.
Returns:
LayoutTensor - A newly allocated tensor with the specified layout
view
view[address_space: AddressSpace](self, ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space]) -> LayoutTensor[dtype, __layout, MutableAnyOrigin, address_space=address_space]
Creates a tensor view over existing memory.
Note: Fails to compile if layout is not set, address spaces don't match, or tensor is circular.
Parameters:
- address_space (
AddressSpace): Memory address space for the tensor (generic, shared, local).
Args:
- ptr (
UnsafePointer[SIMD[dtype, 1], address_space=address_space]): Pointer to memory region to create the view over.
Returns:
LayoutTensor - A tensor view over the specified memory region with the current layout.
circular
circular(self) -> LayoutTensorBuild[dtype, __layout=__layout, __layout_init=__layout_init, __address_space=__address_space, __circular=True]
Enables circular indexing for the tensor.
Returns:
LayoutTensorBuild - A new builder with circular indexing enabled.
iter
iter(self, ptr: UnsafePointer[SIMD[dtype, 1], address_space=__address_space], bound: Int) -> LayoutTensorIter[dtype, __layout, MutableAnyOrigin, address_space=__address_space, circular=__circular]
Creates an iterator over tensor elements.
Note: Fails to compile if layout is not set or dimensions are not known.
Args:
- ptr (
UnsafePointer[SIMD[dtype, 1], address_space=__address_space]): Pointer to memory region. - bound (
Int): Upper bound for iteration.
Returns:
LayoutTensorIter - An iterator over tensor elements.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!