From 1f96eb31a0afd2cc20eb462a19a5b1434498fac1 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Fri, 18 Jul 2025 10:43:30 -0700 Subject: [PATCH 01/11] init --- backends/apple/coreml/compiler/coreml_preprocess.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py index bf390698705..85e0ccb03d9 100644 --- a/backends/apple/coreml/compiler/coreml_preprocess.py +++ b/backends/apple/coreml/compiler/coreml_preprocess.py @@ -282,9 +282,9 @@ def get_model_debug_info(model_package_dir: Path) -> Optional[ModelDebugInfo]: if delegate_info is None: return None - debug_handle_to_operation_path_mapping: Optional[Dict[str, Any]] = ( - delegate_info.get("mapping", None) - ) + debug_handle_to_operation_path_mapping: Optional[ + Dict[str, Any] + ] = delegate_info.get("mapping", None) if debug_handle_to_operation_path_mapping is None: return None From bd7fe58d64063a15e7f93e52bd7b587b5716a911 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Fri, 18 Jul 2025 13:15:28 -0700 Subject: [PATCH 02/11] up --- backends/apple/coreml/compiler/torch_ops.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py index 18a840972c6..b05633bb898 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/backends/apple/coreml/compiler/torch_ops.py @@ -30,7 +30,6 @@ def transpose_copy(context, node): transpose(context, node) - # https://github.com/apple/coremltools/pull/2557 @register_torch_op(override=False) def unbind_copy(context, node): From 97a66ed2366557c661b489076a7d10745dc868b2 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Fri, 18 Jul 2025 13:15:44 -0700 Subject: [PATCH 03/11] up --- backends/apple/coreml/compiler/torch_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py index b05633bb898..18a840972c6 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/backends/apple/coreml/compiler/torch_ops.py @@ -30,6 +30,7 @@ def transpose_copy(context, node): transpose(context, node) + # https://github.com/apple/coremltools/pull/2557 @register_torch_op(override=False) def unbind_copy(context, node): From c103181a7e2d92135d0ba37872789452b8e8a1d1 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Mon, 21 Jul 2025 09:41:27 -0700 Subject: [PATCH 04/11] lint --- backends/apple/coreml/compiler/coreml_preprocess.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py index 85e0ccb03d9..bf390698705 100644 --- a/backends/apple/coreml/compiler/coreml_preprocess.py +++ b/backends/apple/coreml/compiler/coreml_preprocess.py @@ -282,9 +282,9 @@ def get_model_debug_info(model_package_dir: Path) -> Optional[ModelDebugInfo]: if delegate_info is None: return None - debug_handle_to_operation_path_mapping: Optional[ - Dict[str, Any] - ] = delegate_info.get("mapping", None) + debug_handle_to_operation_path_mapping: Optional[Dict[str, Any]] = ( + delegate_info.get("mapping", None) + ) if debug_handle_to_operation_path_mapping is None: return None From 6d98f7a1edb78da7348808b5b57760f90c066b1c Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Mon, 21 Jul 2025 09:38:51 -0700 Subject: [PATCH 05/11] up --- .github/workflows/_unittest.yml | 2 +- backends/apple/coreml/test/test_torch_ops.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/_unittest.yml b/.github/workflows/_unittest.yml index 63f5e6693b7..5be04de1ccf 100644 --- a/.github/workflows/_unittest.yml +++ b/.github/workflows/_unittest.yml @@ -43,7 +43,7 @@ jobs: macos: uses: pytorch/test-infra/.github/workflows/macos_job.yml@main with: - runner: macos-m1-stable + runner: macos-15 python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py index 323f76afd1b..4d26c892a3e 100644 --- a/backends/apple/coreml/test/test_torch_ops.py +++ b/backends/apple/coreml/test/test_torch_ops.py @@ -2,7 +2,10 @@ # # Please refer to the license found in the LICENSE file in the root directory of the source tree. +<<<<<<< HEAD import platform +======= +>>>>>>> 921dafe9b (up) import sys import unittest From 530a21ce321e3db9feb70c820fc0a03989043559 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Mon, 21 Jul 2025 10:49:26 -0700 Subject: [PATCH 06/11] up --- .github/workflows/_unittest.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/_unittest.yml b/.github/workflows/_unittest.yml index 5be04de1ccf..9657a6d82c5 100644 --- a/.github/workflows/_unittest.yml +++ b/.github/workflows/_unittest.yml @@ -43,7 +43,7 @@ jobs: macos: uses: pytorch/test-infra/.github/workflows/macos_job.yml@main with: - runner: macos-15 + runner: macos-15-xlarge python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} From 217fc7c10fb9090b24f95f20087a80cb236c24fe Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 16 Jul 2025 14:35:07 -0700 Subject: [PATCH 07/11] init --- exir/program/_program.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/exir/program/_program.py b/exir/program/_program.py index 19c06da23bd..4b12352a081 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1076,6 +1076,26 @@ def keep(op): return list(filter(keep, preserve_ops)) +def _can_skip_using_EDGE_DO_NOT_DECOMP( + partitioner: Dict[str, List[Partitioner]], aten_programs: Dict[str, ExportedProgram] +) -> bool: + # THe current design of using EDGE_DO_NOT_DECOMP to prevent decomposition + # has long standing issues. _remove_invalid_ops_for_not_decompose was a band-aid to + # fix some of the issues, but more issues are coming up over time, including a new issue with SDPA + # and contiguous views: https://fb.workplace.com/groups/pytorch.edge.users/permalink/1796069037930048/ + # EDGE_DO_NOT_DECOMP is only needed by partitioners that specify check_op_support + # As a temp fix, we give a more reliable path for backends that do not specify check_op_support + can_skip_using_EDGE_DO_NOT_DECOMP = True + 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) + ) + if check_op_support is not None: + can_skip_using_EDGE_DO_NOT_DECOMP = False + return can_skip_using_EDGE_DO_NOT_DECOMP + def _gen_edge_manager_for_partitioners( partitioner: Dict[str, List[Partitioner]], aten_programs: Dict[str, ExportedProgram], From 027942adc5de94f91d03efe909b7b20f44e83596 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Thu, 17 Jul 2025 17:49:17 -0700 Subject: [PATCH 08/11] updates --- exir/program/_program.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/exir/program/_program.py b/exir/program/_program.py index 4b12352a081..557b271854a 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1096,6 +1096,7 @@ def _can_skip_using_EDGE_DO_NOT_DECOMP( can_skip_using_EDGE_DO_NOT_DECOMP = False return can_skip_using_EDGE_DO_NOT_DECOMP + def _gen_edge_manager_for_partitioners( partitioner: Dict[str, List[Partitioner]], aten_programs: Dict[str, ExportedProgram], @@ -1118,22 +1119,43 @@ def _gen_edge_manager_for_partitioners( 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() 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 + 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 = 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 From 4f7fa468216969d43c27dcd378e0877a52f330cb Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Thu, 17 Jul 2025 19:33:43 -0700 Subject: [PATCH 09/11] up --- exir/program/_program.py | 80 ++++++++++++++++++++++++---------------- 1 file changed, 49 insertions(+), 31 deletions(-) diff --git a/exir/program/_program.py b/exir/program/_program.py index 557b271854a..1f3997e71e3 100644 --- a/exir/program/_program.py +++ b/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, From b5e24849733c3851660ff22914d8239295315b40 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Sun, 20 Jul 2025 18:04:47 -0700 Subject: [PATCH 10/11] up --- exir/program/_program.py | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/exir/program/_program.py b/exir/program/_program.py index 1f3997e71e3..1fc6d6e7224 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1098,22 +1098,6 @@ def _can_skip_using_EDGE_DO_NOT_DECOMP( 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], @@ -1166,9 +1150,9 @@ def _gen_edge_manager_for_partitioners( # 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 if can_skip_using_EDGE_DO_NOT_DECOMP: - ops_set_to_not_decompose_by_program[ - name - ] = all_ops_no_decomp_needing_preservation + 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) @@ -1179,8 +1163,6 @@ def _gen_edge_manager_for_partitioners( 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, From 2579c8526ddf96e47c364d6d467d61cee4bbbf99 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Tue, 22 Jul 2025 12:55:20 -0700 Subject: [PATCH 11/11] up --- .github/workflows/_unittest.yml | 2 +- backends/apple/coreml/test/test_torch_ops.py | 3 --- exir/program/_program.py | 3 +++ 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/_unittest.yml b/.github/workflows/_unittest.yml index 9657a6d82c5..63f5e6693b7 100644 --- a/.github/workflows/_unittest.yml +++ b/.github/workflows/_unittest.yml @@ -43,7 +43,7 @@ jobs: macos: uses: pytorch/test-infra/.github/workflows/macos_job.yml@main with: - runner: macos-15-xlarge + runner: macos-m1-stable python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py index 4d26c892a3e..323f76afd1b 100644 --- a/backends/apple/coreml/test/test_torch_ops.py +++ b/backends/apple/coreml/test/test_torch_ops.py @@ -2,10 +2,7 @@ # # Please refer to the license found in the LICENSE file in the root directory of the source tree. -<<<<<<< HEAD import platform -======= ->>>>>>> 921dafe9b (up) import sys import unittest diff --git a/exir/program/_program.py b/exir/program/_program.py index 1fc6d6e7224..8bbe0833b85 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1123,6 +1123,9 @@ def _gen_edge_manager_for_partitioners( ops_set_to_not_decompose_by_program = {} edge_programs: Dict[str, ExportedProgram] = {} for name, program in aten_programs.items(): + # Functionalize program before asking partitioners to preserve ops + program = program.run_decompositions({}) + if partitioner is not None: # preserve all ops listed by all partitioners first all_ops_no_decomp = set()