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 trait

TileLoader

DRAM→LDS DMA loader contract for tile_rows × tile_cols half-tiles.

Implementations cooperate as a warp group to fill an SMEM half-tile via buffer_load_*_lds. The kernel walks coords in (m_offset, k_offset) GEMM-space; the loader translates them to physical addresses internally. Conformers must be TrivialRegisterPassable so the kernel can pass them by value through closures.

Two conformers ship today:

  • TileLoaderLDS: linear 2D source. Used by matmul A/B operands and by conv's B (filter) operand. The address math is addr = (m_offset * stride) + k_offset.
  • TileLoaderLDSIm2col: NHWC + in-line im2col. Used by conv's A (input) operand. The address math decomposes m_offset → (n, h_out, w_out) and k_offset → (kh, kw, c) at load time; conv geometry (R, S, H, W, stride, dilation, pad) is loader-internal state.

The kernel doesn't have to know which loader is in use; it just advances (m_offset, k_offset) through the K-loop. That's the point of the trait: the conv body and matmul body can share everything except which loader they instantiate.

Implemented traits​

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

dtype​

comptime dtype

tile_cols​

comptime tile_cols

tile_rows​

comptime tile_rows

Required methods​

load_tile​

def load_tile(self, dst: TileTensor[Self.dtype, Engine=dst.Engine, address_space=AddressSpace.SHARED], m_offset: Int, k_offset: Int)

Loads a half-tile from global memory into the SMEM dst.

Issues num_iterations buffer_load_*_lds bursts (per lane) that together fill the tile_rows × tile_cols SMEM half-tile. Each iteration costs one vmcnt-tracked outstanding load per lane: the 4-wave software pipeline relies on this exact accounting.

Args:

Was this page helpful?