`benchmarks.SaturnNPU.scripts.trace_layer_decomposition`

Source: benchmarks/SaturnNPU/scripts/trace_layer_decomposition.py

`benchmarks.SaturnNPU.scripts.trace_layer_decomposition`

Trace PyTorch layer decomposition through all MLIR compilation levels.

Uses MLIR Python bindings to walk the IR at each level and map ops back to their PyTorch origins. Produces a detailed per-layer trace showing:

PyTorch layer → Torch-MLIR ops → Linalg/Input ops → Global-Opt ops

For each op: name, shape, body pattern (for generics), and PyTorch role.

Usage

python scripts/trace_layer_decomposition.py --linalg-input .../phases/module.1.input.mlir --global-opt .../phases/module.4.global-optimization.mlir --output-dir benchmarks/SaturnNPU/

`build_layer_trace(ops, level_name)`

Build a per-layer decomposition from the classified op sequence.

Returns a list of layer dicts, each with the ops assigned to it.

`classify_generic(op)`

Classify a linalg.generic op by inspecting its region body ops.

`parse_mlir_ops(path)`

Parse an MLIR file and return classified op sequence for main$async body.

`write_trace_report(input_layers, gopt_layers, output_dir)`

Write the decomposition trace as markdown + JSON.