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_sram_async_ldmatrix

def copy_dram_to_sram_async_ldmatrix[thread_layout: Layout[thread_layout.shape_types, thread_layout.stride_types], *, simd_size: Int, swizzle: Bool, masked: Bool = False, num_threads: Int = thread_layout.size()](dst: TileTensor[Engine=dst.Engine, address_space=AddressSpace.SHARED, linear_idx_type=dst.linear_idx_type], src: TileTensor[Engine=src.Engine, address_space=src.address_space, linear_idx_type=src.linear_idx_type], num_valid_rows: Int = Int(0))

Copies a tile from global to shared memory, one simd_size vector at a time, optionally swizzled for ldmatrix.

On NVIDIA this is an asynchronous cp.async copy; the caller commits and waits. Elsewhere it falls back to the synchronous copy, where swizzle and masked do not apply.

With swizzle, the destination is swizzled by the pattern make_ldmatrix_swizzle derives from its row stride, so the tile can be read by TensorCore.load_a/load_b with that same row stride. A masked copy zero-fills the vectors of rows at or beyond num_valid_rows.

Parameters:

  • ​thread_layout (Layout[thread_layout.shape_types, thread_layout.stride_types]): Layout of the threads over the copy.
  • ​simd_size (Int): Elements per copied vector.
  • ​swizzle (Bool): Whether to swizzle the destination for ldmatrix.
  • ​masked (Bool): Whether to zero-fill rows past num_valid_rows.
  • ​num_threads (Int): Total threads in the block; those beyond thread_layout.size() do not copy.

Args:

Was this page helpful?