From 243536414ccca92a0ec31ee90462becf50d66e71 Mon Sep 17 00:00:00 2001
From: Zingo Andersen <zingo.andersen@arm.com>
Date: Tue, 8 Oct 2024 14:03:58 -0700
Subject: [PATCH] Arm backend: Show delegation info from aot_arm_compiler
 (#5981)

Summary:
This shows what operators are delegated and not.

Pull Request resolved: https://github.com/pytorch/executorch/pull/5981

Reviewed By: mergennachin

Differential Revision: D64048404

Pulled By: digantdesai

fbshipit-source-id: 8c79c4384e7df771bc7b0fd041c76625b8d87552
---
 examples/arm/aot_arm_compiler.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/examples/arm/aot_arm_compiler.py b/examples/arm/aot_arm_compiler.py
index fed8b8b0b4..ef13a3c346 100644
--- a/examples/arm/aot_arm_compiler.py
+++ b/examples/arm/aot_arm_compiler.py
@@ -10,6 +10,7 @@
 import argparse
 import logging
 import os
+from typing import Optional
 
 import torch
 
@@ -19,8 +20,11 @@
     ArmQuantizer,
     get_symmetric_quantization_config,
 )
+
+from executorch.devtools.backend_debug import get_delegation_info
 from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
 from executorch.extension.export_util.utils import export_to_edge, save_pte_program
+from tabulate import tabulate
 
 # Quantize model if required using the standard export quantizaion flow.
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
@@ -198,6 +202,21 @@ def get_compile_spec(target: str, intermediates: bool) -> ArmCompileSpecBuilder:
     return spec_builder.build()
 
 
+def dump_delegation_info(edge, intermediate_files_folder: Optional[str] = None):
+    graph_module = edge.exported_program().graph_module
+    delegation_info = get_delegation_info(graph_module)
+    df = delegation_info.get_operator_delegation_dataframe()
+    table = tabulate(df, headers="keys", tablefmt="fancy_grid")
+    delegation_info_string = f"Delegation info:\n{delegation_info.get_summary()}\nDelegation table:\n{table}\n"
+    logging.info(delegation_info_string)
+    if intermediate_files_folder is not None:
+        delegation_file_path = os.path.join(
+            intermediate_files_folder, "delegation_info.txt"
+        )
+        with open(delegation_file_path, "w") as file:
+            file.write(delegation_info_string)
+
+
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -313,6 +332,9 @@ def get_args():
     logging.debug(f"Exported graph:\n{edge.exported_program().graph}")
     if args.delegate is True:
         edge = edge.to_backend(ArmPartitioner(compile_spec))
+
+        dump_delegation_info(edge, args.intermediates)
+
         logging.debug(f"Lowered graph:\n{edge.exported_program().graph}")
 
     try: