Mojo struct
EpilogueConfig
@register_passable(trivial)
struct EpilogueConfig[MMA_M: Int, MMA_N: Int, stageN: Int, cta_group: Int, transpose_c: Bool]
Configuration for epilogue stage computations.
Computes the number of stages and other parameters needed for the output epilogue based on MMA and CTA configuration.
Template Parameters: MMA_M: MMA M dimension. MMA_N: MMA N dimension. stageN: Stage width in elements. cta_group: Number of CTAs cooperating (1 or 2). transpose_c: Whether output is transposed.
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
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
bits
comptime bits = 256
cg1_num_stages
comptime cg1_num_stages = (MMA_N // stageN)
cg2_num_stages
comptime cg2_num_stages = (MMA_N // stageN) if (MMA_M == 256) else ((MMA_N // stageN) // 2)
data_paths
comptime data_paths = 16
fragment_size
comptime fragment_size = 4
is_lower_frag_required
comptime is_lower_frag_required = (MMA_M == 64) if (cta_group == 1) else (cta_group == 1).__bool__().__invert__()
num_stages
comptime num_stages = (MMA_N // stageN) if (MMA_M == 256) else ((MMA_N // stageN) // 2) if (cta_group == 2) else EpilogueConfig[MMA_M, MMA_N, stageN, cta_group, transpose_c].cg1_num_stages
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!