Arm Backend: add .pte file visualization in Model Explorer (#16290)

### Summary

Modify `examples/arm/run.sh` and `examples/arm/visualize.py` to be able
to trigger visualization (pte or tosa) of the model.
For the Ethos-U delegate, we need to pass down intermediate .tosa
artifacts. Invoke the PTE adapter API to do that.

### Test plan

I tested with
1. `./examples/arm/run.sh --model_name=mv2 --model_explorer
--visualize_pte`
2. `./examples/arm/run.sh --model_name=mv2 --target="vgf"
--model_explorer --visualize_pte`
3. `./examples/arm/run.sh --model_name=mv2 --model_explorer
--visualize_tosa`


Signed-off-by: Tomás González <[email protected]>
Co-authored-by: Tomás González <[email protected]>

Author: tgonzalezorlandoarm

backends/arm/requirements-arm-ethos-u.txt

@@ -3,4 +3,5 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-ethos-u-vela == 4.4.1
+ethos-u-vela == 4.4.1
+pte-adapter-model-explorer == 0.0.2

examples/arm/run.sh

@@ -43,6 +43,8 @@ select_ops_list="aten::_softmax.out"
 qdq_fusion_op=false
 model_explorer=false
 perf_overlay=false
+visualize_tosa=false
+visualize_pte=false
 
 function help() {
     echo "Usage: $(basename $0) [options]"
@@ -73,8 +75,11 @@ function help() {
     echo "  --et_build_root=<FOLDER>               Executorch build output root folder to use, defaults to ${et_build_root}"
     echo "  --scratch-dir=<FOLDER>                 Path to your Arm scrach dir if you not using default ${arm_scratch_dir}"
     echo "  --qdq_fusion_op                        Enable QDQ fusion op"
-    echo "  --model_explorer                       Enable model explorer to visualize TOSA graph."
-    echo "  --perf_overlay                         With --model_explorer, include performance data from FVP PMU trace."
+    echo "  --model_explorer                       Enable model explorer to visualize a TOSA or PTE model graph."
+    echo "  --visualize_pte                        With --model_explorer, visualize PTE flatbuffer model and delegates. Cannot be used with --visualize_tosa"
+    echo "                                            NOTE:  If PTE contains an Ethos-U delegate, the Ethos-U subgraph will be visualized if aot_arm_compiler_flags is set with the -i flag to include intermediate tosa files."
+    echo "  --visualize_tosa                       With --model_explorer, visualize TOSA flatbuffer model. Cannot be used with --visualize_pte"
+    echo "  --perf_overlay                         With --model_explorer and --visualize_tosa, include performance data from FVP PMU trace."
     exit 0
 }
 
@@ -105,6 +110,8 @@ for arg in "$@"; do
       --qdq_fusion_op) qdq_fusion_op=true;;
       --model_explorer) model_explorer=true ;;
       --perf_overlay) perf_overlay=true ;;
+      --visualize_tosa) visualize_tosa=true ;;
+      --visualize_pte) visualize_pte=true ;;
       *)
       ;;
     esac
@@ -351,12 +358,21 @@ for i in "${!test_model[@]}"; do
     fi
 
     if [ "$model_explorer" = true ]; then
-        tosa_flatbuffer_path=$(find ${output_folder} -name "*TOSA*.tosa" | head -n 1)
         perf_flags=""
         if [ "$perf_overlay" = true ]; then
             perf_flags+="--trace ${output_folder}/pmu_trace.gz --tables ${output_folder}/output/out_debug.xml"
         fi
-        python3 ${script_dir}/visualize.py --model_path ${tosa_flatbuffer_path} ${perf_flags}
+
+        visualization_file=""
+        if [ "$visualize_tosa" = true ]; then
+            visualization_file+=" --tosa"
+        fi
+        if [ "$visualize_pte" = true ]; then
+            visualization_file+=" --pte"
+        fi
+
+        me_flags="${visualization_file} ${perf_flags}"
+        python3 ${script_dir}/visualize.py --model_dir ${output_folder} ${me_flags}
     fi
 done
 

examples/arm/visualize.py

@@ -18,6 +18,7 @@
     visualize_model_explorer,
 )
 from model_explorer import config as model_explorer_config, node_data_builder as ndb
+from model_explorer.config import ModelSource
 
 COMPILER_OP_ID = "scheduled_id"
 
@@ -252,17 +253,53 @@ def validate_perf_mode_args(trace: str, tables: str) -> None:
         )
 
 
+def set_pte_model_explorer_config(model_file, tosa_files, config):
+    from pte_adapter_model_explorer.main import PTEAdapter
+
+    pte_adapter = PTEAdapter()
+
+    settings = {"delegate_file_paths": [str(path) for path in tosa_files]}
+    me_graphs = pte_adapter.convert(model_path=str(model_file), settings=settings)
+
+    # Convert the given model to model explorer graphs.
+
+    graphs_index = len(config.graphs_list)
+    config.graphs_list.append(me_graphs)
+
+    # Construct model source.
+    #
+    # The model source has a special format, in the form of:
+    # graphs://{name}/{graphs_index}
+    model_name = model_file.stem
+    model_source: ModelSource = {"url": f"graphs://{model_name}/{graphs_index}"}
+    config.model_sources.append(model_source)
+
+
+def set_tosa_model_explorer_config(model_file, config):
+    config.add_model_from_path(str(model_file))
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         description="Visualize a model using model explorer."
     )
     parser.add_argument(
-        "--model_path", required=True, type=str, help="Path to the model file"
+        "--model_dir", required=True, type=str, help="Path to the model directory"
+    )
+    parser.add_argument(
+        "--pte",
+        action="store_true",
+        help="Visualize PTE flatbuffer model and delegates. Cannot be used with --tosa",
+    )
+    parser.add_argument(
+        "--tosa",
+        action="store_true",
+        help="Visualize TOSA flatbuffer model. Cannot be used with --pte",
     )
     parser.add_argument(
         "--trace",
         required=False,
-        help="(perf mode) PMU trace JSON.gz file with performance data",
+        help="(perf mode) PMU trace JSON.gz file with performance data. Can only be used together with --tosa",
     )
     parser.add_argument(
         "--tables",
@@ -271,10 +308,37 @@ def main() -> None:
     )
 
     args = parser.parse_args()
-    model_file = Path(args.model_path).resolve()
+    if args.pte and args.tosa:
+        raise ValueError("Cannot use both --pte and --tosa options together")
+
+    model_dir = Path(args.model_dir).resolve()
+
+    tosa_files = list(model_dir.glob("*TOSA*.tosa"))
+
+    model_files = None
+    if args.pte:
+        model_files = list(model_dir.glob("*.pte"))
+    elif args.tosa:
+        model_files = tosa_files
+
+    if not model_files:
+        raise FileNotFoundError(
+            f"No model files found in {model_dir} for the specified format."
+        )
+
+    model_file = model_files[0]
+
     validate_file_exists(model_file)
 
-    config = model_explorer_config().add_model_from_path(str(model_file))
+    config = model_explorer_config()
+
+    extensions = []
+
+    if args.pte:
+        set_pte_model_explorer_config(model_file, tosa_files, config)
+    elif args.tosa:
+        set_tosa_model_explorer_config(model_file, config)
+        extensions.append("tosa_adapter_model_explorer")
 
     if args.trace or args.tables:
         validate_perf_mode_args(args.trace, args.tables)
@@ -287,7 +351,7 @@ def main() -> None:
             "Duration (Cycles)", build_overlay_data(trace_file, tables_file)
         )
 
-    visualize_model_explorer(config=config, extensions=["tosa_adapter_model_explorer"])
+    visualize_model_explorer(config=config, extensions=extensions)
 
 
 if __name__ == "__main__":