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 function

copy_local_to_shared

def copy_local_to_shared[thread_layout: Layout, swizzle: Optional[Swizzle] = None, num_threads: Int = thread_layout.size(), thread_scope: ThreadScope = ThreadScope.BLOCK, block_dim_count: Int = Int(1), *, row_major: Bool = False](dst: LayoutTensor[address_space=AddressSpace.SHARED, element_layout=dst.element_layout, layout_int_type=dst.layout_int_type, linear_idx_type=dst.linear_idx_type, masked=dst.masked, alignment=dst.alignment], src: LayoutTensor[address_space=AddressSpace.LOCAL, element_layout=src.element_layout, layout_int_type=src.layout_int_type, linear_idx_type=src.linear_idx_type, masked=src.masked, alignment=src.alignment])

Synchronously copy data from local memory (registers) to SRAM (shared memory).

This function performs a synchronous copy operation from register memory to shared memory in a GPU context, distributing the workload across multiple threads for parallel execution. It's particularly useful for transferring processed data from registers to shared memory for inter-thread communication.

Performance:

Distributes the copy workload across multiple threads for parallel execution.
Can use swizzling to optimize memory access patterns and reduce bank conflicts.
Optimized for transferring data from registers to shared memory.
On AMD GPUs, the row_major parameter can be used to match the memory access pattern used during prefetching from DRAM to registers.

Notes:

The destination tensor must be in SHARED address space (SRAM).
The source tensor must be in LOCAL address space (registers).
This function is particularly useful in GPU kernels for sharing processed data between threads in the same block.
The row_major parameter is specifically designed for AMD GPUs when using a prefetching pattern from DRAM to SRAM via registers.

Constraints:

Destination tensor must be in SHARED address space.
Source tensor must be in LOCAL address space.
For optimal performance, the thread layout should match the memory access patterns of the tensors.

Parameters:

thread_layout (Layout): Layout defining how threads are organized for the operation. This determines how the workload is distributed among threads.
swizzle (Optional[Swizzle]): Optional swizzling function to rearrange the destination indices, which can improve memory access patterns and reduce bank conflicts.
num_threads (Int): Total number of threads in the thread block. Threads beyond thread_layout.size() will be disabled and not participate in the copy operation.
thread_scope (ThreadScope): Defines whether operations are performed at BLOCK or WARP level. BLOCK scope involves all threads in a thread block, while WARP scope restricts operations to threads within the same warp. Defaults to ThreadScope.BLOCK.
block_dim_count (Int): The number of dimensions in the thread block.
row_major (Bool): Whether to use row-major ordering for the copy operation. This is particularly relevant when prefetching from DRAM to SRAM via registers on AMD GPUs. Defaults to False.

Args:

dst (LayoutTensor[address_space=AddressSpace.SHARED, element_layout=dst.element_layout, layout_int_type=dst.layout_int_type, linear_idx_type=dst.linear_idx_type, masked=dst.masked, alignment=dst.alignment]): The destination tensor, which must be in shared memory (SRAM).
src (LayoutTensor[address_space=AddressSpace.LOCAL, element_layout=src.element_layout, layout_int_type=src.layout_int_type, linear_idx_type=src.linear_idx_type, masked=src.masked, alignment=src.alignment]): The source tensor, which must be in local memory (registers).