Mojo function

smem_subtile

smem_subtile[tile_rows: Int, tile_cols: Int, BN: Int, BK: Int, dtype: DType](smem_ptr: UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED], tile_row: Int, tile_col: Int) -> TileTensor[dtype, Layout[ComptimeInt[tile_rows], ComptimeInt[tile_cols], ComptimeInt[tile_cols], ComptimeInt[1]], MutAnyOrigin, address_space=AddressSpace.SHARED]

Creates a flat TileTensor sub-view of a blocked SMEM layout.

The blocked layout has num_repeats contiguous BN×BK blocks. This function computes the physical offset for a block-aligned tile and returns a row-major TileTensor view with strides (tile_cols, 1).

Correct only when tile_cols == BK (tiles don't cross block boundaries in the column dimension).

Parameters:

tile_rows (Int): Height of the sub-tile.
tile_cols (Int): Width of the sub-tile (must equal BK for block alignment).
BN (Int): Number of rows per block (full block height).
BK (Int): Number of columns per block (full block width).
dtype (DType): Element data type.

Args:

smem_ptr (UnsafePointer): Base pointer to the SMEM allocation.
tile_row (Int): Tile row index (0-based, in units of tile_rows).
tile_col (Int): Tile column index (0-based, in units of tile_cols).

Returns:

TileTensor: A TileTensor view into the specified sub-tile region.