migrate etrecord generation after to_edge_transform_and_lower to new infra (pytorch#13153)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#13058 by
@Gasoonjia
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/34/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/34/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/gasoonjia/34/orig
@diff-train-skip-merge

Differential Revision: D79420502

Co-authored-by: gasoonjia <[email protected]>
Co-authored-by: Gasoonjia <[email protected]>

Author: 2 people

backends/qualcomm/tests/utils.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 import collections
-import copy
 import os
 import subprocess
 import tempfile
@@ -30,7 +29,7 @@
     get_soc_to_chipset_map,
     to_edge_transform_and_lower_to_qnn,
 )
-from executorch.devtools import generate_etrecord, Inspector
+from executorch.devtools import Inspector
 from executorch.devtools.inspector._inspector_utils import TimeScale
 from executorch.examples.qualcomm.utils import (
     generate_inputs,
@@ -512,11 +511,9 @@ def lower_module_and_test_output(
             skip_node_id_set=skip_node_id_set,
             skip_node_op_set=skip_node_op_set,
             skip_mutable_buffer=skip_mutable_buffer,
+            generate_etrecord=self.enable_profile,
         )
 
-        # this is needed for the ETRecord as lowering modifies the graph in-place
-        edge_copy = copy.deepcopy(delegated_program)
-
         exec_prog = delegated_program.to_executorch(
             exir.ExecutorchBackendConfig(
                 # For shared buffer, user must pass the memory address
@@ -543,7 +540,7 @@ def lower_module_and_test_output(
 
         etrecord_path = "etrecord.bin"
         if self.enable_profile:
-            generate_etrecord(etrecord_path, edge_copy, exec_prog)
+            exec_prog.get_etrecord().save(etrecord_path)
         # Check numerics
         if (
             assert_output_equal

backends/qualcomm/utils/utils.py

@@ -334,6 +334,7 @@ def to_edge_transform_and_lower_to_qnn(
     skip_node_id_set: Optional[set] = None,
     skip_node_op_set: Optional[set] = None,
     skip_mutable_buffer: bool = False,
+    generate_etrecord: bool = False,
 ) -> EdgeProgramManager:
     """
     Transforms and lowers a given PyTorch module to the QNN backend.
@@ -442,6 +443,7 @@ def ensure_graph_specific_dict(value, graph_names):
         partitioner=qnn_partitioners,
         constant_methods=constant_methods,
         compile_config=qnn_edge_config(),
+        generate_etrecord=generate_etrecord,
     )
 
 

examples/apple/coreml/scripts/export.py

@@ -223,7 +223,6 @@ def main():
     pte_base_name = get_pte_base_name(args)
     if args.use_partitioner:
         model = model.eval()
-        assert not args.generate_etrecord, "ETRecord is not supported with partitioner"
         ep = torch.export.export(
             model,
             args=example_args,
@@ -234,9 +233,12 @@ def main():
         delegated_program = exir.to_edge_transform_and_lower(
             ep,
             partitioner=[CoreMLPartitioner(compile_specs=compile_specs)],
+            generate_etrecord=args.generate_etrecord,
         )
         exec_program = delegated_program.to_executorch()
         save_pte_program(exec_program, pte_base_name)
+        if args.generate_etrecord:
+            exec_program.get_etrecord().save(f"{pte_base_name}_coreml_etrecord.bin")
         if args.run_with_pybindings:
             run_with_pybindings(
                 executorch_program=exec_program,

examples/qualcomm/scripts/export_example.py

@@ -1,6 +1,5 @@
 # pyre-ignore-all-errors
 import argparse
-import copy
 
 import torch
 from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer
@@ -10,7 +9,6 @@
     get_soc_to_chipset_map,
     to_edge_transform_and_lower_to_qnn,
 )
-from executorch.devtools import generate_etrecord
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
 from executorch.exir.capture._config import ExecutorchBackendConfig
@@ -107,19 +105,16 @@ def main() -> None:
         backend_options=backend_options,
     )
     delegated_program = to_edge_transform_and_lower_to_qnn(
-        m, example_inputs, compile_spec
+        m, example_inputs, compile_spec, generate_etrecord=args.generate_etrecord
     )
 
-    # this is needed for the ETRecord as lowering modifies the graph in-place
-    edge_copy = copy.deepcopy(delegated_program)
-
     executorch_program = delegated_program.to_executorch(
         config=ExecutorchBackendConfig(extract_delegate_segments=False)
     )
 
     if args.generate_etrecord:
         etrecord_path = args.output_folder + "etrecord.bin"
-        generate_etrecord(etrecord_path, edge_copy, executorch_program)
+        executorch_program.get_etrecord().save(etrecord_path)
 
     save_pte_program(executorch_program, args.model_name, args.output_folder)
 

examples/qualcomm/util_scripts/gen_etrecord.py

@@ -1,4 +1,3 @@
-import copy
 import os
 
 import torch
@@ -10,7 +9,7 @@
     QcomChipset,
     to_edge_transform_and_lower_to_qnn,
 )
-from executorch.devtools import generate_etrecord, Inspector
+from executorch.devtools import Inspector
 from executorch.devtools.inspector._inspector_utils import TimeScale
 from executorch.examples.qualcomm.utils import (
     make_quantizer,
@@ -46,11 +45,9 @@ def main(args):
         module=converted,
         inputs=sample_input,
         compiler_specs=compiler_specs,
+        generate_etrecord=True,
     )
 
-    # for inspector API
-    edge_copy = copy.deepcopy(edge_prog_mgr)
-
     # store pte file
     exec_prog = edge_prog_mgr.to_executorch()
     with open(f"{pte_filename}.pte", "wb") as f:
@@ -71,7 +68,7 @@ def main(args):
 
     # pull etdump back and display the statistics
     adb.pull_etdump(".")
-    generate_etrecord("etrecord.bin", edge_copy, exec_prog)
+    exec_prog.get_etrecord().save("etrecord.bin")
     inspector = Inspector(
         etdump_path="etdump.etdp",
         etrecord="etrecord.bin",

examples/xnnpack/aot_compiler.py

@@ -9,12 +9,10 @@
 # Example script for exporting simple models to flatbuffer
 
 import argparse
-import copy
 import logging
 
 import torch
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
-from executorch.devtools import generate_etrecord
 from executorch.exir import (
     EdgeCompileConfig,
     ExecutorchBackendConfig,
@@ -60,6 +58,7 @@
         "-r",
         "--etrecord",
         required=False,
+        default="",
         help="Generate and save an ETRecord to the given file location",
     )
     parser.add_argument("-o", "--output_dir", default=".", help="output directory")
@@ -103,18 +102,16 @@
             _check_ir_validity=False if args.quantize else True,
             _skip_dim_order=True,  # TODO(T182187531): enable dim order in xnnpack
         ),
+        generate_etrecord=args.etrecord,
     )
     logging.info(f"Exported and lowered graph:\n{edge.exported_program().graph}")
 
-    # this is needed for the ETRecord as lowering modifies the graph in-place
-    edge_copy = copy.deepcopy(edge)
-
     exec_prog = edge.to_executorch(
         config=ExecutorchBackendConfig(extract_delegate_segments=False)
     )
 
-    if args.etrecord is not None:
-        generate_etrecord(args.etrecord, edge_copy, exec_prog)
+    if args.etrecord:
+        exec_prog.get_etrecord().save(args.etrecord)
         logging.info(f"Saved ETRecord to {args.etrecord}")
 
     quant_tag = "q8" if args.quantize else "fp32"