Commit 5db1bfb
committed
[ExecuTorch][Export][1/N] Export API pipeline re-architecture, making it composable
RFC: #12660
This diff introduces composable architecture for the export pipeline.
Changes:
1. Introduces notion of `PipelineArtifact` which encompasses stage artifacts and run context necessary to execute a stage and the pipeline.
1. Move all stages to new file to better architect the component
1. Export api will now accept `pipeline_stages` and honors the sequence if provided otherwise fallback to default sequence: `source_transform -> quantize -> torch.export -> to_edge_transform_and_lower -> to_executorch`
1. Validate if pipeline provided is valid, error out if not.
1. Add unittests to test pipeline sequence and stages
With this one can execute a partial pipeline, for example, one can just do `TORCH_EXPORT -> TO_EDGE_TRANSFORM_AND_LOWER -> TO_EXECUTORCH`
Current limitation:
1. `TORCH_EXPORT` stage is mandatory to avoid someone passing incorrect input such as `ExportedProgram` instead of nn.module, this enforcement will anchor the pipeline to have expected stages to run.
- This limitation will be removed if we choose to add the support to take `ExportedProgram` as input.
Note:
1. this diff is large because I moved stages to separate file and added bunch of unittests but core functionality that is added can be reviewed in `export.py` file. CC: digantdesai
2. Export component was supposed to mimic stages design in tester and now it is.
Fixes: #12928
Differential Revision: [D79120574](https://our.internmc.facebook.com/intern/diff/D79120574/)
ghstack-source-id: 299106076
Pull Request resolved: #129361 parent 45846c8 commit 5db1bfb
File tree
7 files changed
+946
-916
lines changed- export
- tests
7 files changed
+946
-916
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
15 | 15 | | |
16 | 16 | | |
17 | 17 | | |
18 | | - | |
19 | 18 | | |
20 | 19 | | |
21 | 20 | | |
| |||
31 | 30 | | |
32 | 31 | | |
33 | 32 | | |
| 33 | + | |
34 | 34 | | |
35 | 35 | | |
36 | 36 | | |
37 | 37 | | |
38 | 38 | | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| 59 | + | |
39 | 60 | | |
40 | 61 | | |
41 | 62 | | |
| |||
48 | 69 | | |
49 | 70 | | |
50 | 71 | | |
| 72 | + | |
51 | 73 | | |
52 | 74 | | |
53 | 75 | | |
| |||
0 commit comments