Mojo struct
TMEMFragment
@register_passable(trivial)
struct TMEMFragment[accum_type: DType, epilogue_type: DType, frag_size: Int]
Accumulator fragment pair from tensor memory.
SM100 TMEM stores data in upper/lower fragment pairs due to the physical layout of tensor memory datapaths.
Template Parameters: accum_type: Accumulator data type (e.g., float32). epilogue_type: Epilogue data type after casting (e.g., bfloat16). frag_size: Number of elements per fragment.
Fields
- upper (
SIMD[accum_type, frag_size]): - lower (
SIMD[accum_type, frag_size]): - has_lower (
Bool):
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
Methods
__init__
__init__(has_lower: Bool = True) -> Self
Initialize empty fragments.
Args:
- has_lower (
Bool): Whether lower fragment is needed (based on MMA config).
cast_upper
cast_upper(self) -> SIMD[epilogue_type, frag_size]
Cast upper fragment to epilogue type.
Returns:
SIMD: Upper fragment cast to epilogue_type.
cast_lower
cast_lower(self) -> SIMD[epilogue_type, frag_size]
Cast lower fragment to epilogue type.
Returns:
SIMD: Lower fragment cast to epilogue_type.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!