Add access to edge_program in ArmPassManager (pytorch#5542)

Summary:
Change-Id: Iecb1af0001dee7c48961129f0d49644cddfc18a0

Pull Request resolved: pytorch#5542

Reviewed By: cccclai

Differential Revision: D63262789

Pulled By: digantdesai

fbshipit-source-id: 8219a97aeb3ee25e318724297a39f86987557600

Author: Erik-Lundell

backends/arm/arm_backend.py

@@ -217,7 +217,7 @@ def preprocess(  # noqa: C901
         # const data directly. Path created and data written only in debug builds.
         tosa_graph = ts.TosaSerializer(artifact_path)
         graph_module = ArmPassManager().transform_to_backend_pipeline(
-            graph_module=edge_program.graph_module, compile_spec=compile_spec
+            exported_program=edge_program, compile_spec=compile_spec
         )
 
         node_visitors = get_node_visitors(edge_program)

backends/arm/passes/arm_pass_manager.py

@@ -22,6 +22,7 @@
 )
 from executorch.backends.arm.passes.remove_clone_pass import RemoveClonePass
 from executorch.backends.arm.passes.size_adjust_conv2d_pass import SizeAdjustConv2DPass
+from executorch.exir import ExportedProgram
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.pass_manager import PassManager
 
@@ -32,7 +33,7 @@ def _transform(self, graph_module: torch.fx.GraphModule):
         return self(graph_module).graph_module
 
     def transform_to_backend_pipeline(
-        self, graph_module: torch.fx.GraphModule, compile_spec: list[CompileSpec]
+        self, exported_program: ExportedProgram, compile_spec: list[CompileSpec]
     ):
         """Apply passes before transforming program to backend"""
         self.add_pass(SizeAdjustConv2DPass())
@@ -46,4 +47,4 @@ def transform_to_backend_pipeline(
                 if memory_format == "nhwc":
                     self.add_pass(AnnotateChannelsLastDimOrder())
 
-        return self._transform(graph_module)
+        return self._transform(exported_program.graph_module)