Mojo struct

LocalToSharedTileCopier

struct LocalToSharedTileCopier[thread_layout: Layout[thread_layout.shape_types, thread_layout.stride_types], *, swizzle: Optional[Swizzle] = None, num_threads: Int = thread_layout.size(), thread_scope: ThreadScope = ThreadScope.BLOCK]

A TileCopier that moves a tile from registers into shared memory.

The AMD row_major prefetch pattern and fp32 -> half precision downcast are not supported.

Parameters

thread_layout (Layout): Layout describing how threads are organized over the copy.
swizzle (Optional): Optional swizzle applied to the shared-memory destination; the same swizzle must be used by any subsequent reader of the tile.
num_threads (Int): Total number of threads in the thread block. Threads beyond thread_layout.size() do not participate.
thread_scope (ThreadScope): Scope at which thread operations are performed.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, TileCopier

`comptime` members

`dst_address_space`

comptime dst_address_space = AddressSpace.SHARED

Destination AddressSpace this copier writes to.

`src_address_space`

comptime src_address_space = AddressSpace.LOCAL

Source AddressSpace this copier reads from.

Methods

`copy`

copy[element_size: Int](self, dst: TileTensor[dst.dtype, dst.LayoutType, dst.origin, address_space=LocalToSharedTileCopier[thread_layout, swizzle=swizzle, num_threads=num_threads, thread_scope=thread_scope].dst_address_space, linear_idx_type=dst.linear_idx_type, element_size=element_size], src: TileTensor[src.dtype, src.LayoutType, src.origin, address_space=LocalToSharedTileCopier[thread_layout, swizzle=swizzle, num_threads=num_threads, thread_scope=thread_scope].src_address_space, linear_idx_type=src.linear_idx_type, element_size=element_size])

Copies src in local memory into dst in shared memory.

Parameters:

element_size (Int): Number of scalar elements per logical element.

Args:

dst (TileTensor): Destination tile in shared memory.
src (TileTensor): Source tile in local memory.

Parameters​

Implemented traits​

comptime members​

dst_address_space​

src_address_space​

Methods​

copy​