NXP backend: Add visualization support to aot_neutron_compile.py.

Author: MartinPavella

examples/nxp/aot_neutron_compile.py

@@ -21,6 +21,9 @@
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
 from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
+from executorch.devtools.visualization.visualization_utils import (
+    visualize_with_clusters,
+)
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
 from executorch.exir import (
@@ -210,6 +213,12 @@ def _get_batch_size(data):
         nargs="*",
         help="List of operators not to delegate. E.g., --operators_not_to_delegate aten::convolution aten::mm",
     )
+    parser.add_argument(
+        "--visualize",
+        choices=["show", "store"],
+        help="Visualize the lowered program. `show` launches a browser tab with the visualization. `store` stores the "
+        "visualization in a json file for later inspection. See `docs/source/visualize-with-clusters.md` for details.",
+    )
 
     args = parser.parse_args()
 
@@ -237,7 +246,7 @@ def _get_batch_size(data):
         module = post_training_quantize(module, calibration_inputs)
 
     if args.so_library is not None:
-        logging.debug(f"Loading libraries: {args.so_library} and {args.portable_lib}")
+        logging.debug(f"Loading libraries: {args.so_library}")
         torch.ops.load_library(args.so_library)
 
     if args.test:
@@ -284,7 +293,7 @@ def _get_batch_size(data):
             raise RuntimeError(
                 e.args[0]
                 + ".\nThis likely due to an external so library not being loaded. Supply a path to it with the "
-                "--portable_lib flag."
+                "--so_library flag."
             ).with_traceback(e.__traceback__) from None
         else:
             raise e
@@ -301,3 +310,13 @@ def executorch_program_to_str(ep, verbose=False):
         "_nxp_delegate" if args.delegate is True else ""
     )
     save_pte_program(exec_prog, model_name)
+
+    # 7. Optionally visualize the model.
+    if args.visualize == "show":
+        visualize_with_clusters(exec_prog.exported_program())
+    elif args.visualize == "store":
+        file_name = f"{args.model_name}-visualization.json"
+        logging.info(
+            f"Saved the graph visualization in `{file_name}`. It can be opened using the ModelExplorer."
+        )
+        visualize_with_clusters(exec_prog.exported_program(), file_name)

examples/nxp/setup.sh

@@ -8,3 +8,9 @@ set -u
 
 # Install neutron-converter
 pip install --index-url https://eiq.nxp.com/repository neutron_converter_SDK_25_09
+
+# Get the directory of the current script
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Install the required visualization dependencies.
+"${SCRIPT_DIR}/../../devtools/install_requirements.sh"