Skip to main content

Mojo struct

StandardTilePayload

@register_passable(trivial) struct StandardTilePayload[a_type: DType, b_type: DType, a_tile_layout: Layout, b_tile_layout: Layout, num_pipeline_stages: Int]

Tile payload for standard matmul (A and B tiles).

Fields

  • a_tiles (StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].ATileArray):
  • b_tiles (StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].BTileArray):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, TilePayload

comptime members

__copyinit__is_trivial

comptime __copyinit__is_trivial = True

__del__is_trivial

comptime __del__is_trivial = True

__moveinit__is_trivial

comptime __moveinit__is_trivial = True

ATile

comptime ATile = StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].ATileArray.Tile

ATileArray

comptime ATileArray = SMemTileArrayType[a_type, a_tile_layout, num_pipeline_stages, 128]

BTile

comptime BTile = StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].BTileArray.Tile

BTileArray

comptime BTileArray = SMemTileArrayType[b_type, b_tile_layout, num_pipeline_stages, 128]

Methods

__init__

__init__(a_tiles: SMemTileArrayType[a_type, a_tile_layout, num_pipeline_stages, 128], b_tiles: SMemTileArrayType[b_type, b_tile_layout, num_pipeline_stages, 128]) -> Self

get_tile

get_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> Tuple[StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].ATile, StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].BTile]

Get A and B tiles at the specified stage and k-group index.

Returns:

Tuple

get_a_tile

get_a_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].ATile

Get A tile at the specified stage and k-group index.

Returns:

StandardTilePayload

get_b_tile

get_b_tile[k_group_size: Int](self, stage: UInt32, k_idx: Int) -> StandardTilePayload[a_type, b_type, a_tile_layout, b_tile_layout, num_pipeline_stages].BTile

Get B tile at the specified stage and k-group index.

Returns:

StandardTilePayload

Was this page helpful?