Skip to content

samples.research.model_computation_graph_generation.mlir_flow_parser

Source: samples/research/model_computation_graph_generation/mlir_flow_parser.py

samples.research.model_computation_graph_generation.mlir_flow_parser

MLIR Stream Parser - Extract dispatch dependencies and build DAG Parses IREE MLIR stream files to create a dependency graph of dispatch operations. Handles concurrent execution blocks to identify parallel dispatches.

DispatchNode dataclass

Represents a dispatch operation in the MLIR stream

MLIRStreamParser

Parse MLIR stream files and extract dispatch dependencies with concurrent execution info

build_graph()

Build NetworkX DAG from parsed dispatches

export_to_json(output_path)

Export graph structure to JSON with concurrent execution info

get_concurrent_groups()

Get all concurrent groups of dispatches

get_critical_path()

Get the critical path (longest path) in the DAG

get_max_parallelism()

Get maximum number of dispatches that can run in parallel

get_parallel_stages()

Get dispatch stages that can be executed in parallel

load_file()

Load MLIR file content

parse()

Main parsing pipeline

parse_dependencies()

Parse the stream command execution to extract dispatch dependencies and concurrent groups

parse_dispatches()

Extract all dispatch operations from the MLIR file

print_summary()

Print summary of parsed information

visualize(output_path=None)

Visualize the dependency graph with concurrent groups highlighted

main()

Main entry point