Mojo struct
BlackwellWarpProfilingWorkspaceManager
@register_passable(trivial)
struct BlackwellWarpProfilingWorkspaceManager[load_warps: UInt32, mma_warps: UInt32, scheduler_warps: UInt32, epilogue_warps: UInt32, max_entries_per_warp: UInt32]
This struct manages the profiling workspace. The workspaces consists of equal sized chunks, the total number of which is equal to the total number of active SMs. Each SM chunk consists of sequences of entries, with a maximum number of entries per warp role.
Template Parameters: load_warps: Number of warps specialized for load operations mma_warps: Number of warps specialized for matrix multiply-accumulate operations scheduler_warps: Number of warps specialized for scheduling operations epilogue_warps: Number of warps specialized for epilogue operations max_entries_per_warp: Maximum number of entries per warp (common across all warp roles)
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
entries_per_sm
comptime entries_per_sm = max_entries_per_warp.__rmul__[DType.uint32, 1](BlackwellWarpProfilingWorkspaceManager[load_warps, mma_warps, scheduler_warps, epilogue_warps, max_entries_per_warp].total_warp_roles)
header
comptime header = "time_start,time_end,sm_id,block_idx_x,block_idx_y,role,entry_idx\n"
sm_count
comptime sm_count = B200.sm_count
total_data_points
comptime total_data_points = 7
total_warp_roles
comptime total_warp_roles = 4
Methods
get_workspace
write_to_workspace
static write_to_workspace[warp_role: UInt32](sm_idx: UInt32, entry_idx: UInt32, workspace: Span[UInt64, MutAnyOrigin], timeline: Tuple[UInt64, UInt64])
dump_workspace_as_csv
static dump_workspace_as_csv(ctx: DeviceContext, workspace: Span[UInt64, MutAnyOrigin], filename: StringSlice[StaticConstantOrigin])
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!