Update on "Remove usages of 'to_edge_with_preserve_ops'"

Replace with to_edge.

Differential Revision: [D78311705](https://our.internmc.facebook.com/intern/diff/D78311705/)

[ghstack-poisoned]

Author: lucylq

backends/cadence/aot/compiler.py

@@ -242,7 +242,7 @@ def _lower_ep_to_edge(
             # Allow specific non-core aten ops in the IR.
             _core_aten_ops_exception_list=TO_EDGE_OP_EXCEPTION_LIST
             + (core_aten_exceptions or []),
-            _preserve_ops=TO_EDGE_PRESERVE_OPS,
+            preserve_ops=TO_EDGE_PRESERVE_OPS,
         ),
         constant_methods=constant_methods,
     )

examples/apple/coreml/llama/export.py

@@ -201,7 +201,7 @@ def main() -> None:
         compile_config=EdgeCompileConfig(
             _check_ir_validity=False,
             _skip_dim_order=True,
-            _preserve_ops=[
+            preserve_ops=[
                 torch.ops.aten.scaled_dot_product_attention.default,
                 # preserve norm op for numerical stability
                 torch.ops.aten.linalg_vector_norm.default,

exir/capture/_config.py

@@ -39,6 +39,8 @@ class EdgeCompileConfig:
     _check_ir_validity: bool = True
     # TODO(larryliu): remove this
     _use_edge_ops: bool = True
+    # TODO(gasoonjia): remove this
+    _skip_dim_order: bool = False
     # Allow core ATen ops check to be skipped for certain ops, but continue with the rest of the checks.
     # Note: only use this for core ATen ops that are missing decompositions. This is temporary,
     # enabling verification on the rest of the program until decomposition coverage is improved.
@@ -47,9 +49,7 @@ class EdgeCompileConfig:
     )
     # Allow ops to be preserved in the graph, i.e., prevent them from being decomposed.
     # These may be core or non-core ATen ops; custom ops should not be here.
-    _preserve_ops: List[torch.torch._ops.OpOverload] = field(default_factory=list)
-    # TODO(gasoonjia): remove this
-    _skip_dim_order: bool = False
+    preserve_ops: List[torch.torch._ops.OpOverload] = field(default_factory=list)
 
 
 @compatibility(is_backward_compatible=False)

exir/program/_program.py

@@ -1382,8 +1382,8 @@ def to_edge(
         table = _default_decomposition_table()
         preserve_ops = []
         if compile_config:
-            preserve_ops = compile_config._preserve_ops
-            for op in compile_config._preserve_ops:
+            preserve_ops = compile_config.preserve_ops
+            for op in compile_config.preserve_ops:
                 table.pop(op, None)
         program = program.run_decompositions(table)
         edge_programs[name] = _generate_edge_program(

exir/program/test/test_program.py

@@ -784,7 +784,7 @@ def _test_to_edge_with_preserved_ops(
         self, program, preserved_ops, expected_preserved_ops
     ):
         edge = to_edge(
-            program, compile_config=EdgeCompileConfig(_preserve_ops=preserved_ops)
+            program, compile_config=EdgeCompileConfig(preserve_ops=preserved_ops)
         )
 
         def count_nodes(graph_module, target):

exir/verification/test/test_verifier.py

@@ -171,7 +171,7 @@ def forward(self, x):
                 return x.expand(2, 2, 2, 2)
 
         model = TestExpand()
-        config = EdgeCompileConfig(_preserve_ops=[torch.ops.aten.expand.default])
+        config = EdgeCompileConfig(preserve_ops=[torch.ops.aten.expand.default])
         export_model = export(model, (torch.randn(2, 2, 2, 2),), strict=True)
         with self.assertRaises(RuntimeError):
             to_edge(export_model, compile_config=config)

exir/verification/verifier.py

@@ -98,8 +98,8 @@ def EXIRATenDialectVerifier(  # noqa: C901
             _core_aten_ops_exception_list.extend(
                 edge_compile_config._core_aten_ops_exception_list
             )
-        if edge_compile_config._preserve_ops:
-            _preserve_ops.extend(edge_compile_config._preserve_ops)
+        if edge_compile_config.preserve_ops:
+            _preserve_ops.extend(edge_compile_config.preserve_ops)
 
     class _EXIRATenDialectVerifier(EXIRATenDialectVerifierBase):
         dialect = "OLD_EXIR_ATEN"
@@ -181,7 +181,7 @@ def get_aten_verifier(config: EdgeCompileConfig):
         EXIRATenDialectVerifier(
             class_only=True,
             core_aten_ops_exception_list=config._core_aten_ops_exception_list,
-            preserve_ops=config._preserve_ops,
+            preserve_ops=config.preserve_ops,
         )
         if config._check_ir_validity
         else EXIRATenDialectVerifierBase
@@ -253,8 +253,8 @@ def EXIREdgeDialectVerifier(  # noqa: C901
             _core_aten_ops_exception_list.extend(
                 edge_compile_config._core_aten_ops_exception_list
             )
-        if edge_compile_config._preserve_ops:
-            _preserve_ops.extend(edge_compile_config._preserve_ops)
+        if edge_compile_config.preserve_ops:
+            _preserve_ops.extend(edge_compile_config.preserve_ops)
 
     class _EXIREdgeDialectVerifier(Verifier):
         dialect = "EDGE"