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_dram

def copy_local_to_dram[dst_thread_layout: Layout, num_threads: Int = dst_thread_layout.size(), thread_scope: ThreadScope = ThreadScope.BLOCK, block_dim_count: Int = Int(1)](dst: LayoutTensor[address_space=dst.address_space, 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=src.address_space, 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])

Efficiently copy data from registers (LOCAL) to global memory (DRAM).

This function implements a high-performance memory transfer operation from register memory to global memory. It distributes the copy operation across multiple threads for maximum throughput while handling bounds checking for safety.

Constraints:

• The source tensor must be in LOCAL address space (registers).
• The destination tensor must be in GENERIC or GLOBAL address space (DRAM).
• Both tensors must have compatible data types.

Parameters:

• ​dst_thread_layout (Layout): The layout used to distribute the destination tensor across threads. This determines how the workload is divided among participating threads.
• ​num_threads (Int): Total number of threads in the thread block. Threads beyond dst_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.

Args:

• ​dst (LayoutTensor[address_space=dst.address_space, 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 in global memory (DRAM).
• ​src (LayoutTensor[address_space=src.address_space, 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 in register memory (LOCAL) to be copied.

def copy_local_to_dram[dst_thread_layout: Layout, num_threads: Int = dst_thread_layout.size(), thread_scope: ThreadScope = ThreadScope.BLOCK, block_dim_count: Int = Int(1)](dst: LayoutTensor[address_space=dst.address_space, 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=src.address_space, 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], dst_base: LayoutTensor[address_space=dst_base.address_space, element_layout=dst_base.element_layout, layout_int_type=dst_base.layout_int_type, linear_idx_type=dst_base.linear_idx_type, masked=dst_base.masked, alignment=dst_base.alignment])

Efficiently copy data from registers (LOCAL) to global memory (DRAM) on AMD GPUs.

This function implements an optimized memory transfer operation specifically for AMD GPU architectures. It utilizes the hardware's buffer_store intrinsic to efficiently transfer data from registers to global memory while handling bounds checking. The function distributes the copy operation across multiple threads for maximum throughput.

Notes:

• This function is particularly useful for writing computed results from registers back to global memory with minimal latency.
• The offset calculation is optimized for performance rather than flexibility.

Constraints:

• Only supported on AMD GPUs.
• Destination tensor must be in GLOBAL address space.
• Source tensor must be in LOCAL address space.
• Data types must match between source and destination tensors.

Parameters:

• ​dst_thread_layout (Layout): The layout used to distribute the destination tensor across threads. This determines how the workload is divided among participating threads.
• ​num_threads (Int): Total number of threads in the thread block. Threads beyond dst_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.

Args:

• ​dst (LayoutTensor[address_space=dst.address_space, 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 in global memory (DRAM).
• ​src (LayoutTensor[address_space=src.address_space, 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 in register memory (LOCAL address space) to be copied.
• ​dst_base (LayoutTensor[address_space=dst_base.address_space, element_layout=dst_base.element_layout, layout_int_type=dst_base.layout_int_type, linear_idx_type=dst_base.linear_idx_type, masked=dst_base.masked, alignment=dst_base.alignment]): The original global memory tensor from which dst is derived. This is used to construct the buffer descriptor required by AMD's buffer_store intrinsic.