For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
TileWriterThreadwise
struct TileWriterThreadwise[dtype: DType, dst_layout: TensorLayout, dst_origin: MutOrigin, dst_storage: TensorStorage, dst_linear_idx_type: DType, //, thread_layout: Layout[thread_layout.shape_types, thread_layout.stride_types], simd_size: Int, half_tile: Bool = False, swapAB: Bool = False]
Writes shared-memory tiles to global memory using per-thread vectorized stores.
Implements SMemTileWriter without hardware TMA: each thread reads a
SIMD-width chunk from a swizzled shared-memory tile and writes it directly
to the destination global tensor. Supports an optional A/B swap mapping and
a half-tile mode for the x2 masked-consumer case.
Parameters
- dtype (
DType): Data type of the source and destination tiles (inferred). - dst_layout (
TensorLayout): Layout of the destination global tensor (inferred). - dst_origin (
MutOrigin): Origin type of the destination global tensor (inferred). - dst_storage (
TensorStorage): Storage type of the destination global tensor (inferred). - dst_linear_idx_type (
DType): Linear index type for the destination tensor (inferred). - thread_layout (
Layout[thread_layout.shape_types, thread_layout.stride_types]): Layout mapping threads across the tile for vectorized stores. - simd_size (
Int): SIMD vector width, in elements, used for vectorized stores. - half_tile (
Bool): Whether to write only half the tile for the masked x2 consumer case (defaults to False). - swapAB (
Bool): Whether to transpose the A and B matrix mapping (defaults to False).
Fields
- dst (
TileWriterThreadwise[thread_layout, simd_size, half_tile, swapAB].DstType): - thread_idx (
Int):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
SMemTileWriter,
TrivialRegisterPassable
comptime members
DstType
comptime DstType = TileTensor[dtype, dst_layout, dst_origin, Storage=dst_storage, linear_idx_type=dst_linear_idx_type]
Methods
__init__
def __init__(dst: TileTensor[dtype, dst_layout, dst_origin, Storage=dst_storage, linear_idx_type=dst_linear_idx_type], thread_idx: Int) -> Self
Initialize the threadwise tile writer.
Args:
- dst (
TileTensor[dtype, dst_layout, dst_origin, Storage=dst_storage, linear_idx_type=dst_linear_idx_type]): Destination tensor in global memory. - thread_idx (
Int): Thread index within the consumer warp group.
write_tile
def write_tile(self, src: TileTensor[dtype, Storage=src.Storage, address_space=AddressSpace.SHARED, linear_idx_type=src.linear_idx_type], coords: Tuple[Int, Int])
Write a tile using thread-distributed stores.
Each thread writes a portion of the tile with proper swizzling for optimal memory access patterns.
Args:
- src (
TileTensor[dtype, Storage=src.Storage, address_space=AddressSpace.SHARED, linear_idx_type=src.linear_idx_type]): Source tile in shared memory. - coords (
Tuple[Int, Int]): Tile indices (row_tile, col_tile) in the destination matrix.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!