fix etrecord lost after to_backend

Differential Revision: D79823973

Pull Request resolved: #13197

Author: Gasoonjia

devtools/etrecord/tests/TARGETS

@@ -7,12 +7,7 @@ python_unittest(
     name = "etrecord_test",
     srcs = ["etrecord_test.py"],
     deps = [
-        "//caffe2:torch",
-        "//executorch/devtools/bundled_program:config",
-        "//executorch/devtools/bundled_program:core",
-        "//executorch/devtools/etrecord:etrecord",
-        "//executorch/exir:lib",
-        "//executorch/exir/tests:models",
+        ":etrecord_test_library"
     ],
 )
 
@@ -26,5 +21,6 @@ python_library(
         "//executorch/devtools/etrecord:etrecord",
         "//executorch/exir:lib",
         "//executorch/exir/tests:models",
+        "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
     ],
 )

devtools/etrecord/tests/etrecord_test.py

@@ -15,6 +15,7 @@
 import executorch.exir.tests.models as models
 import torch
 from executorch import exir
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite
 from executorch.devtools.bundled_program.core import BundledProgram
 from executorch.devtools.etrecord import generate_etrecord, parse_etrecord
@@ -369,6 +370,41 @@ def test_get_etrecord_from_executorch_program_manager(self):
         self.assertEqual(etrecord._debug_handle_map, et_manager.debug_handle_map)
         self.assertEqual(etrecord._delegate_map, et_manager.delegate_map)
 
+    def test_get_etrecord_from_executorch_program_manager_with_partitioner(self):
+        """Test getting ETRecord from ExecutorchProgramManager using get_etrecord() method."""
+        f = models.BasicSinMax()
+        aten_program = export(f, f.get_random_inputs(), strict=True)
+
+        # Generate edge manager with ETRecord
+        edge_manager = to_edge_transform_and_lower(
+            aten_program,
+            partitioner=[XnnpackPartitioner()],
+            generate_etrecord=True,
+        )
+
+        # Convert to executorch
+        et_manager = edge_manager.to_executorch()
+
+        # Test get_etrecord method
+        etrecord = et_manager.get_etrecord()
+        self.assertIsNotNone(etrecord)
+        self.assert_etrecord_saveable(etrecord)
+
+        # Verify the data matches the original input
+        self.check_graph_closeness(
+            etrecord.exported_program,
+            aten_program.graph_module,
+        )
+        self.assertEqual(
+            etrecord.export_graph_id,
+            id(aten_program.graph),
+        )
+
+        # Verify the executorch program data matches
+        # ETRecord stores data directly (not JSON serialized), so compare with original data
+        self.assertEqual(etrecord._debug_handle_map, et_manager.debug_handle_map)
+        self.assertEqual(etrecord._delegate_map, et_manager.delegate_map)
+
     def test_get_etrecord_from_executorch_program_manager_without_generation(self):
         """Test getting ETRecord from ExecutorchProgramManager when ETRecord was not generated."""
         f = models.BasicSinMax()
@@ -400,6 +436,7 @@ def test_to_edge_transform_and_lower_etrecord_save_and_parse(self):
         # Generate edge manager with ETRecord
         edge_manager = to_edge_transform_and_lower(
             aten_program,
+            partitioner=[XnnpackPartitioner()],
             generate_etrecord=True,
         )
 

exir/program/_program.py

@@ -1583,10 +1583,13 @@ def transform(
                 new_programs[name].graph_module
             )
 
-        return EdgeProgramManager(
+        epm = EdgeProgramManager(
             new_programs, copy.deepcopy(self._config_methods), compile_config
         )
 
+        epm._etrecord = self._etrecord
+        return epm
+
     @et_logger("to_backend")
     def to_backend(
         self,
@@ -1629,12 +1632,15 @@ def to_backend(
 
         new_edge_programs = to_backend(method_to_programs_and_partitioners)
         config = EdgeCompileConfig(_check_ir_validity=False)
-        return EdgeProgramManager(
+        epm = EdgeProgramManager(
             new_edge_programs,
             copy.deepcopy(self._config_methods),
             config,
         )
 
+        epm._etrecord = self._etrecord
+        return epm
+
     @et_logger("to_executorch")
     def to_executorch(
         self,