set example input to only the very first partition

Author: Gasoonjia

backends/aoti/aoti_backend.py

@@ -75,8 +75,6 @@ def preprocess(
 
         named_data_store = NamedDataStore()
 
-        # copy_edge_program = copy.deepcopy(edge_program)
-
         # Move the edge_program from CPU to CUDA for aoti compile
         cuda_edge_program = move_to_device_pass(edge_program, "cuda")
 

exir/backend/backend_api.py

@@ -268,7 +268,9 @@ def _partition_and_lower_one_graph_module(
     """
     Partitioned and lowered the graph module based on the partition tag, this is to handle one graph module.
     """
-    for tag, delegation_spec in partition_result.partition_tags.items():
+    for idx, (tag, delegation_spec) in enumerate(
+        partition_result.partition_tags.items()
+    ):
         # Create partition with nodes containing this tag. There should only be
         # one contained submodule per tag
         node_list = _get_node_list_with_same_tag(
@@ -311,6 +313,7 @@ def _partition_and_lower_one_graph_module(
             tag,
             call_module_node,
             is_submodule,
+            idx == 0,
         )
 
         lowered_submodule = to_backend(
@@ -452,7 +455,9 @@ def _create_partitions_in_graph_module(
     is_submodule: bool,
 ) -> Dict[str, List[torch.fx.Node]]:
     backend_id_to_submodule_name = {}
-    for tag, delegation_spec in partition_result.partition_tags.items():
+    for idx, (tag, delegation_spec) in enumerate(
+        partition_result.partition_tags.items()
+    ):
         # Create partition with nodes containing this tag. There should only be
         # one contained submodule per tag
         node_list = _get_node_list_with_same_tag(
@@ -492,6 +497,7 @@ def _create_partitions_in_graph_module(
             tag,
             call_module_node,
             is_submodule,
+            idx == 0,
         )
         call_module_node.meta["backend_id"] = delegation_spec.backend_id
         call_module_node.meta["compile_spec"] = delegation_spec.compile_specs

exir/lowered_backend_module.py

@@ -682,6 +682,7 @@ def create_exported_program_from_submodule(
     tag: str,
     call_module_node: torch.fx.Node,
     is_submodule: bool,
+    is_first_partition: bool = False,
 ) -> Tuple[ExportedProgram, Dict[str, InputSpec], Dict[str, OutputSpec]]:
     """
     Creates an ExportedProgram from the given submodule using the parameters and buffers
@@ -720,6 +721,11 @@ def create_exported_program_from_submodule(
     in_spec = pytree.tree_flatten((tuple(subgraph_signature.user_inputs), {}))[1]
     out_spec = pytree.tree_flatten(subgraph_signature.user_outputs)[1]
 
+    # only the example inputs of first parition equals to the example inputs of the owning program
+    submodule_exmaple_inputs = (
+        owning_program.example_inputs if is_first_partition else None
+    )
+
     return (
         ExportedProgram(
             root=submodule,
@@ -735,7 +741,7 @@ def create_exported_program_from_submodule(
                     ),
                 )
             ],
-            example_inputs=owning_program.example_inputs,
+            example_inputs=submodule_exmaple_inputs,
             constants=subgraph_constants,
             verifiers=[owning_program.verifier],
         ),