Skip to main content

Mojo struct

BlockScaledSmem

struct BlockScaledSmem[a_type: DType, b_type: DType, c_type: DType, sfa_dtype: DType, sfb_dtype: DType, transpose_b: Bool, *, config: BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]]

SMEM struct for block-scaled matmul with CLC scheduler pipeline.

Thin wrapper over BlockScaledTileCore + SmemPipelineBundle.

Fields

  • core (BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core):
  • pipelines (BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Pipelines):

Implemented traits

AnyType, ImplicitlyDestructible

comptime members

Core

comptime Core = BlockScaledTileCore[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config]

Pipelines

comptime Pipelines = SmemPipelineBundle[BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.num_group_pipeline_stages, BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.num_accum_pipeline_stages, config.num_clc_pipeline_stages, BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.Payload]

Methods

a_tiles

a_tiles(ref[AddressSpace._value] self) -> BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.ATileArray

Get A tile array accessor.

Returns:

BlockScaledSmem

b_tiles

b_tiles(ref[AddressSpace._value] self) -> BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.BTileArray

Get B tile array accessor.

Returns:

BlockScaledSmem

c_tiles

c_tiles(ref[AddressSpace._value] self) -> BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.CTileArray

Get C tile array accessor.

Returns:

BlockScaledSmem

sfa_tiles

sfa_tiles(ref[AddressSpace._value] self) -> BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.SFATileArray

Get SFA tile array accessor.

Returns:

BlockScaledSmem

sfb_tiles

sfb_tiles(ref[AddressSpace._value] self) -> BlockScaledSmem[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b, config=config].Core.SFBTileArray

Get SFB tile array accessor.

Returns:

BlockScaledSmem

ab_pipeline_size

static ab_pipeline_size() -> Int

Total size of A+B tiles for all pipeline stages (in elements).

Returns:

Int

sf_pipeline_size

static sf_pipeline_size() -> Int

Total size of SFA+SFB tiles for all pipeline stages (in elements).

Returns:

Int

c_output_size

static c_output_size() -> Int

Size of C tiles for all output stages (in elements).

Returns:

Int

total_tile_size

static total_tile_size() -> Int

Total tile storage size (A+B+SFA+SFB+C) in elements.

Returns:

Int

Was this page helpful?