From 28c5fe2550399786f062a8617215e0d7d07f44f1 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 16 Jul 2025 15:29:03 -0700 Subject: [PATCH 1/4] Update base for 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] From f5405a633737967d2f917b24a54187b7397fb322 Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 17 Jul 2025 09:41:38 -0700 Subject: [PATCH 2/4] Remove usages of 'to_edge_with_preserve_ops' Differential Revision: D78311705 Pull Request resolved: https://github.com/pytorch/executorch/pull/12471 --- backends/cadence/aot/compiler.py | 12 ++++++------ examples/apple/coreml/llama/export.py | 16 ++++++++-------- exir/capture/_config.py | 6 +++--- exir/program/_program.py | 4 ++-- exir/program/test/test_program.py | 2 +- exir/verification/test/test_verifier.py | 2 +- exir/verification/verifier.py | 10 +++++----- 7 files changed, 26 insertions(+), 26 deletions(-) diff --git a/backends/cadence/aot/compiler.py b/backends/cadence/aot/compiler.py index bfa07fc03cf..4a4c74da0fd 100644 --- a/backends/cadence/aot/compiler.py +++ b/backends/cadence/aot/compiler.py @@ -34,7 +34,7 @@ ) from executorch.exir.passes import ToOutVarPass from executorch.exir.passes.sym_shape_eval_pass import HintBasedSymShapeEvalPass -from executorch.exir.program._program import to_edge_with_preserved_ops +from executorch.exir.program._program import to_edge from torch._inductor.decomposition import remove_decompositions from torch.export.exported_program import ExportedProgram @@ -219,9 +219,9 @@ def quantize_pt2( torch.ops.aten.angle.default, torch.ops.aten.rms_norm.default, ] -TO_EDGE_PRESERVE_OPS: tuple[torch._ops.OpOverload, ...] = ( +TO_EDGE_PRESERVE_OPS: list[torch._ops.OpOverload, ...] = [ torch.ops.aten.rms_norm.default, -) +] def _lower_ep_to_edge( @@ -233,18 +233,18 @@ def _lower_ep_to_edge( """ Lower an ExportedProgram to an EdgeProgramManager (in edge IR). """ - # Call to_edge_with_preserved_ops to convert the graph to edge IR. + # Call to_edge to convert the graph to edge IR. # Note: dim_order is skipped (https://github.com/pytorch/executorch/issues/3704) - edge_prog_manager = to_edge_with_preserved_ops( + edge_prog_manager = to_edge( expo_program, compile_config=EdgeCompileConfig( _skip_dim_order=True, # 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, ), constant_methods=constant_methods, - preserve_ops=TO_EDGE_PRESERVE_OPS, ) if dump_graphs: diff --git a/examples/apple/coreml/llama/export.py b/examples/apple/coreml/llama/export.py index a367a14c595..b2df12d4abd 100644 --- a/examples/apple/coreml/llama/export.py +++ b/examples/apple/coreml/llama/export.py @@ -27,7 +27,7 @@ from executorch.exir.passes import MemoryPlanningPass from executorch.exir.passes.quant_fusion_pass import QuantFusionPass from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass -from executorch.exir.program._program import to_edge_with_preserved_ops +from executorch.exir.program._program import to_edge from executorch.extension.export_util.utils import save_pte_program @@ -196,17 +196,17 @@ def main() -> None: print("Exported program") print(ep) - edge_manager = to_edge_with_preserved_ops( + edge_manager = to_edge( ep, - preserve_ops=[ - torch.ops.aten.scaled_dot_product_attention.default, - # preserve norm op for numerical stability - torch.ops.aten.linalg_vector_norm.default, - torch.ops.aten.reciprocal.default, - ], compile_config=EdgeCompileConfig( _check_ir_validity=False, _skip_dim_order=True, + preserve_ops=[ + torch.ops.aten.scaled_dot_product_attention.default, + # preserve norm op for numerical stability + torch.ops.aten.linalg_vector_norm.default, + torch.ops.aten.reciprocal.default, + ], ), ) print("Edge program") diff --git a/exir/capture/_config.py b/exir/capture/_config.py index 835bc60dad3..b2252e122c9 100644 --- a/exir/capture/_config.py +++ b/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) diff --git a/exir/program/_program.py b/exir/program/_program.py index 3e06d01788a..61abfa7b9e2 100644 --- a/exir/program/_program.py +++ b/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( diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index 7173b4d50b5..da5647936aa 100644 --- a/exir/program/test/test_program.py +++ b/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): diff --git a/exir/verification/test/test_verifier.py b/exir/verification/test/test_verifier.py index 2be4aeac3ab..79ca7c9e226 100644 --- a/exir/verification/test/test_verifier.py +++ b/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) diff --git a/exir/verification/verifier.py b/exir/verification/verifier.py index ed304e99fc1..6b79b924cd2 100644 --- a/exir/verification/verifier.py +++ b/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" From ac67efcbdc5759f6d967ea99a0461857c0f7ddd3 Mon Sep 17 00:00:00 2001 From: pytorchbot Date: Wed, 16 Jul 2025 13:32:59 -0400 Subject: [PATCH 3/4] Add preserve_ops to EdgeCompileConfig (#12546) 1. Add `preserve_ops` to `EdgeCompileConfig` 2. Remove preserved ops from the decomposition table in `to_edge`. 3. Add checks to the verifier ensuring that preserved ops do not have mutations or views. 4. Update 'core_aten_exception_list' to be 'preserved_ops' in `to_edge_transform_and_lower`. Context/Usage **core_aten_ops_exception_list** - Contains operators that are missing a decomposition to core aten. - Exclude these so that verification can still be run on the rest of the graph. - Ideally, this list should be empty. **preserve_ops** - Contains operators that the user specifically does not want decomposed. - Must be aten; custom ops are ignored by verifier. Edge case: - If an aten operator does not have a decomp, and the user specifically wants it to be preserved, put it in preserve_ops rather than core_aten_ops_exception_list. Differential Revision: [D78298749](https://our.internmc.facebook.com/intern/diff/D78298749/) --- exir/program/test/test_program.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index da5647936aa..217820e4eb2 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -784,7 +784,11 @@ def _test_to_edge_with_preserved_ops( self, program, preserved_ops, expected_preserved_ops ): edge = to_edge( +<<<<<<< HEAD program, compile_config=EdgeCompileConfig(preserve_ops=preserved_ops) +======= + program, compile_config=EdgeCompileConfig(_preserve_ops=preserved_ops) +>>>>>>> 0012ffaf9f (Add preserve_ops to EdgeCompileConfig (#12546)) ) def count_nodes(graph_module, target): From 4b83733747a0f947f6afbd54cc03b840659c86b8 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 16 Jul 2025 15:29:03 -0700 Subject: [PATCH 4/4] Update base for 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] --- exir/program/test/test_program.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index 217820e4eb2..da5647936aa 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -784,11 +784,7 @@ def _test_to_edge_with_preserved_ops( self, program, preserved_ops, expected_preserved_ops ): edge = to_edge( -<<<<<<< HEAD program, compile_config=EdgeCompileConfig(preserve_ops=preserved_ops) -======= - program, compile_config=EdgeCompileConfig(_preserve_ops=preserved_ops) ->>>>>>> 0012ffaf9f (Add preserve_ops to EdgeCompileConfig (#12546)) ) def count_nodes(graph_module, target):