Mojo struct
EpilogueApplier
@register_passable(trivial)
struct EpilogueApplier[MMA_M: Int, stageN: Int, num_stages: Int, repeats: Int, cta_group: Int, transpose_c: Bool]
Apply element-wise epilogue lambda to register fragments.
Fields
- coords (
EpilogueApplier[MMA_M, stageN, num_stages, repeats, cta_group, transpose_c].Coords): - warp_id (
UInt32): - lane_id (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
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
Coords
comptime Coords = FragmentCoords[stageN, repeats]
Methods
__init__
__init__(warp_id: UInt32, lane_id: UInt32) -> Self
compute_staged_coords
compute_staged_coords(self, stage: UInt32, c_row: UInt32, c_col: UInt32) -> Tuple[UInt32, UInt32]
Compute global coords with warp and stage offsets (layout-dependent).
Returns:
apply_to_fragment
apply_to_fragment[epilogue_dtype: DType, frag_size: Int, compute_lambda_fn: elementwise_compute_lambda_type](self, mut frag: SIMD[epilogue_dtype, frag_size], staged_row: UInt32, staged_col: UInt32, is_upper: Bool)
Apply epilogue lambda to fragment elements with global coords.
apply_to_both_fragments
apply_to_both_fragments[epilogue_dtype: DType, frag_size: Int, compute_lambda_fn: elementwise_compute_lambda_type, is_lower_frag_required: Bool](self, mut upper_frag: SIMD[epilogue_dtype, frag_size], mut lower_frag: SIMD[epilogue_dtype, frag_size], stage: UInt32, c_row: UInt32, c_col: UInt32) -> Tuple[SIMD[epilogue_dtype, frag_size], SIMD[epilogue_dtype, frag_size]]
Apply epilogue to both fragments (main entry point).
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!