SIGN IN SIGN UP

Demo CoreML Partitioner (#1981)

Summary:
## TLDR
This PR adds ViT as an e2e example of leveraging CoreML partitioner

💁‍♂️ Tips for reviewers: Review commit by commit would be easier, since commit 1 is just some polishing. It is commit 2 that adds the e2e example

## Details: Motivation and Example Setup
CoreML partitioner is intended for non-resident models to benefit from CoreML as much as possible. Concretely, if a model cannot be fully converted to CoreML, then CoreML partitioner would help to have the convertible part run on CoreML backend. Unfortunately, for the models we have tried, we cannot find a perfect "torch exportable but CoreML unconvertible" example: they are either fully CoreML resident, or torch unexportable.

As a result, here we artificially create an example by utilizing the `skip_ops_for_coreml_delegation`. It is still a realistic example:
1. In many use cases, we may not want the initial (or final) embedding part to run on CoreML, especially if they are small
2. Our intended graph break emerges from this example

## Appendix: Exported Program
Use `CoreMLPartitioner()` (real user usage)
```
        def forward(self, arg152_1: "f32[1, 3, 224, 224]"):
            # No stacktrace found for following nodes
            lowered_module_0 = self.lowered_module_0
            executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0, arg152_1);  lowered_module_0 = arg152_1 = None
            getitem: "f32[1, 1000]" = executorch_call_delegate[0];  executorch_call_delegate = None
            return (getitem,)
```
Use `CoreMLPartitioner(skip_ops_for_coreml_delegation=["aten.convolution.default"])` (for demo purpose)
```
        def forward(self, arg2_1: "f32[768, 3, 16, 16]", arg3_1: "f32[768]", arg152_1: "f32[1, 3, 224, 224]"):
            # File: /Volumes/data/Software/Mine/executorch_121559468/env/lib/python3.10/site-packages/torchvision/models/vision_transformer.py:277 in _process_input, code: x = self.conv_proj(x)
            aten_convolution_default: "f32[1, 768, 14, 14]" = executorch_exir_dialects_edge__ops_aten_convolution_default(arg152_1, arg2_1, arg3_1, [16, 16], [0, 0], [1, 1], False, [0, 0], 1);  arg152_1 = arg2_1 = arg3_1 = None

            # No stacktrace found for following nodes
            lowered_module_0 = self.lowered_module_0
            executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0, aten_convolution_default);  lowered_module_0 = aten_convolution_default = None
            getitem: "f32[1, 1000]" = executorch_call_delegate[0];  executorch_call_delegate = None
            return (getitem,)
```

Pull Request resolved: https://github.com/pytorch/executorch/pull/1981

Reviewed By: cccclai

Differential Revision: D53794713

Pulled By: shoumikhin

fbshipit-source-id: 8d97e7e55fe2f7b787b5ff585e55de5cd391e5c3
Y
yifan_shen3 committed
bdde144cea55ed173b9aa73829d29e15b900fa4a
Parent: 1d07fc4
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 2/15/2024, 7:20:07 AM