hespas.chakra_gen.chakra_trace_generator
Functions
|
|
|
|
|
|
|
Classes
|
- hespas.chakra_gen.chakra_trace_generator.get_block_type(block_type)
- class hespas.chakra_gen.chakra_trace_generator.ChakraTraceGen(config, loaded_dep_graph=None)
Bases:
object- __init__(config, loaded_dep_graph=None)
Initialize the pipeline with a compute estimator, mlir module dependency graph, and configuration. :param compute_estimator: The compute estimator used for execution. :param modules: A data dependent DAG of modules to execute. :param config: Configuration dictionary for execution parameters.
- collective_to_chakra_collective(collective)
- create_chakra_traces(stats_print_filter=None, stats_out_filter=None)
Generate Chakra traces for executed modules.
Generates one trace file per device. When all collectives use a single replica group (all devices together), the traces are identical and hardlinked. When any collective has split replica groups (e.g. [[0,1],[2,3]]), each device gets its own file with the correct pg_name on comm nodes.
- _write_comm_group_json(comm_group_registry)
Write communicator group JSON file for astra-sim.
Maps each comm_group_id to its list of device IDs.
- Args:
comm_group_registry: Dict mapping device-id tuple to group_id.
- get_modules_dep_graph()
- hespas.chakra_gen.chakra_trace_generator.setup_output_dir(output_dir, clean=True)
- hespas.chakra_gen.chakra_trace_generator.create_chakra_traces(config_path, output_dir=None, mlir_file=None, num_threads=-1, clean=False, stats_print_filter=None, stats_out_filter=None, num_npus=None, split_fn=None, block_lim=None, merge=None)
- hespas.chakra_gen.chakra_trace_generator.get_arg_parser()
- hespas.chakra_gen.chakra_trace_generator.main(args=None)