up

Author: metascroy

exir/program/_program.py

@@ -1089,14 +1089,31 @@ def _can_skip_using_EDGE_DO_NOT_DECOMP(
     for name, program in aten_programs.items():
         if partitioner is not None:
             for curr_partitioner in partitioner.get(name, []):
-                curr_ops_no_decomp, check_op_support = (
-                    curr_partitioner.ops_to_not_decompose(program)
-                )
+                (
+                    curr_ops_no_decomp,
+                    check_op_support,
+                ) = curr_partitioner.ops_to_not_decompose(program)
                 if check_op_support is not None:
                     can_skip_using_EDGE_DO_NOT_DECOMP = False
     return can_skip_using_EDGE_DO_NOT_DECOMP
 
 
+def _replace_view_with_view_copy(program: ExportedProgram) -> ExportedProgram:
+    program = program.run_decompositions({})
+    new_gm = ReplaceViewOpsWithViewCopyOpsPass()(program.graph_module).graph_module
+    program = ExportedProgram(
+        root=new_gm,
+        graph=new_gm.graph,
+        graph_signature=_get_updated_graph_signature(program.graph_signature, new_gm),
+        state_dict=program.state_dict,
+        range_constraints=program.range_constraints,
+        module_call_graph=program.module_call_graph,
+        example_inputs=program.example_inputs,
+        constants=program.constants,
+    )
+    return program
+
+
 def _gen_edge_manager_for_partitioners(
     partitioner: Dict[str, List[Partitioner]],
     aten_programs: Dict[str, ExportedProgram],
@@ -1116,58 +1133,55 @@ def _gen_edge_manager_for_partitioners(
           on nodes with preserved aten targets. They are then replaces with transformed ops to
           keep them through the second pass of decompositions
     """
+    can_skip_using_EDGE_DO_NOT_DECOMP = _can_skip_using_EDGE_DO_NOT_DECOMP(
+        partitioner, aten_programs
+    )
     ops_set_to_not_decompose_by_program = {}
     edge_programs: Dict[str, ExportedProgram] = {}
     for name, program in aten_programs.items():
-        # Functionalize program without doing any decompositions
-        program = program.run_decompositions({})
-        ReplaceViewOpsWithViewCopyOpsPass()(program.graph_module)
-
-        print(program)
-
         if partitioner is not None:
             # preserve all ops listed by all partitioners first
             all_ops_no_decomp = set()
+            all_ops_no_decomp_needing_preservation = []
             for curr_partitioner in partitioner.get(name, []):
                 curr_ops_no_decomp, _ = curr_partitioner.ops_to_not_decompose(program)
-<<<<<<< HEAD
-                curr_ops_no_decomp = _remove_invalid_ops_for_not_decompose(
-                    curr_ops_no_decomp
-                )
-=======
->>>>>>> ec44f8478 (updates)
                 all_ops_no_decomp |= set(curr_ops_no_decomp)
-            
+
             # If not using the can_skip_using_EDGE_DO_NOT_DECOMP path, we need to remove invalid ops
-            # Otherwise there will be issues 
+            # Otherwise there will be issues
             if not can_skip_using_EDGE_DO_NOT_DECOMP:
-                all_ops_no_decomp = _remove_invalid_ops_for_not_decompose(list(all_ops_no_decomp))
+                all_ops_no_decomp = _remove_invalid_ops_for_not_decompose(
+                    list(all_ops_no_decomp)
+                )
                 all_ops_no_decomp = set(all_ops_no_decomp)
 
             # Run default decompositions, except for those in all_ops_no_decomp
             table = _default_decomposition_table()
             for op in all_ops_no_decomp:
-<<<<<<< HEAD
-                table.pop(op, None)
-
-=======
                 if table.pop(op, None) is not None:
                     all_ops_no_decomp_needing_preservation.append(op)
->>>>>>> ec44f8478 (updates)
             program = program.run_decompositions(table)
 
             # Among all the preserved aten ops, use the check_op_fn to do an additional
             # check on which ops need to be preserved and which ops need to be decomposed
             # Those which are truly preserved will be replaced with transformed ops
-            ops_set_to_not_decompose_by_program[name] = (
-                _replace_aten_ops_with_transformed_ops(name, program, partitioner) or []
-            )
-        program = program.run_decompositions(_default_decomposition_table())
+            if can_skip_using_EDGE_DO_NOT_DECOMP:
+                ops_set_to_not_decompose_by_program[
+                    name
+                ] = all_ops_no_decomp_needing_preservation
+            else:
+                ops_set_to_not_decompose_by_program[name] = (
+                    _replace_aten_ops_with_transformed_ops(name, program, partitioner)
+                    or []
+                )
 
-        _restore_transformed_ops_to_aten_ops(program)
+        if not can_skip_using_EDGE_DO_NOT_DECOMP:
+            program = program.run_decompositions(_default_decomposition_table())
+            _restore_transformed_ops_to_aten_ops(program)
 
+        # Edge will complain if there are view ops requested for preservation, so we replace them with view_copy
+        program = _replace_view_with_view_copy(program)
         edge_programs[name] = program
-
         edge_programs[name] = _generate_edge_program(
             config,
             program,
@@ -1211,7 +1225,7 @@ def collect_named_data_store_outputs(
 
 
 @et_logger("to_edge_transform_and_lower")
-def to_edge_transform_and_lower(
+def to_edge_transform_and_lower(  # noqa: C901
     programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
     transform_passes: Optional[
         Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
@@ -1276,6 +1290,9 @@ def to_edge_transform_and_lower(
     elif partitioner is None:
         partitioner = {name: [] for name in aten_programs.keys()}
 
+    can_skip_using_EDGE_DO_NOT_DECOMP = _can_skip_using_EDGE_DO_NOT_DECOMP(
+        partitioner, aten_programs
+    )
     edge_manager = _gen_edge_manager_for_partitioners(
         partitioner, aten_programs, config, constant_methods
     )
@@ -1301,7 +1318,8 @@ def to_edge_transform_and_lower(
             curr_op_set, check_op_support = curr_partitioner.ops_to_not_decompose(
                 program
             )
-            curr_op_set = _remove_invalid_ops_for_not_decompose(curr_op_set)
+            if not can_skip_using_EDGE_DO_NOT_DECOMP:
+                curr_op_set = _remove_invalid_ops_for_not_decompose(curr_op_set)
             ops_set_to_not_decompose = ops_set_to_not_decompose.union(curr_op_set)
             _sanity_check_graph_for_non_decomp_ops(
                 name,