IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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_dram_to_local

def copy_dram_to_local[src_thread_layout: Layout, num_threads: Int = src_thread_layout.size(), thread_scope: ThreadScope = ThreadScope.BLOCK, block_dim_count: Int = Int(1), cache_policy: CacheOperation = CacheOperation.ALWAYS](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], src_base: LayoutTensor[address_space=src_base.address_space, element_layout=src_base.element_layout, layout_int_type=src_base.layout_int_type, linear_idx_type=src_base.linear_idx_type, masked=src_base.masked, alignment=src_base.alignment], offset: Optional[Int] = None)

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

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

Notes:

  • The offset calculation method significantly impacts performance. Current implementation optimizes for throughput over flexibility.
  • This function is particularly useful for prefetching data into registers before performing computations, reducing memory access latency.

Constraints:

  • Only supported on AMD GPUs.
  • The destination element layout size must match the SIMD width.
  • Source fragments must be rank 2 with known dimensions.

Parameters:

  • ​src_thread_layout (Layout): The layout used to distribute the source 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 src_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.
  • ​cache_policy (CacheOperation): The cache policy to use for the copy operation. Defaults to CacheOperation.ALWAYS.

Args:

def copy_dram_to_local[src_thread_layout: Layout, num_threads: Int = src_thread_layout.size(), thread_scope: ThreadScope = ThreadScope.BLOCK, block_dim_count: Int = Int(1), cache_policy: CacheOperation = CacheOperation.ALWAYS](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_iter: LayoutTensorIter[address_space=src_iter.address_space, alignment=src_iter.alignment, circular=src_iter.circular, axis=src_iter.axis, layout_int_type=src_iter.layout_int_type, linear_idx_type=src_iter.linear_idx_type, masked=src_iter.masked], bounds: UInt32)

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

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

Notes:

  • The offset calculation method significantly impacts performance. Current implementation optimizes for throughput over flexibility.
  • This function is particularly useful for prefetching data into registers before performing computations, reducing memory access latency.

Constraints:

  • Only supported on AMD GPUs.
  • The destination element layout size must match the SIMD width.
  • Source fragments must be rank 2 with known dimensions.

Parameters:

  • ​src_thread_layout (Layout): The layout used to distribute the source 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 src_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.
  • ​cache_policy (CacheOperation): The cache policy to use for the copy operation. Defaults to CacheOperation.ALWAYS.

Args:

def copy_dram_to_local[src_thread_layout: Layout, num_threads: Int = src_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 global memory (DRAM) to registers.

This function implements an optimized memory transfer operation from global memory to register 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 GLOBAL address space (DRAM).
  • The destination tensor must be in LOCAL address space (registers).
  • Both tensors must have compatible data types.

Parameters:

  • ​src_thread_layout (Layout): The layout used to distribute the source 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 src_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:

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact