Commit adfd1d7
authored
Use the backend's copy hook in single-backend to_backend (pytorch#23339)
## What is wrong today
Before ExecuTorch calls a backend's `preprocess`, it copies the program,
so a backend that edits the graph cannot change the caller's program.
`BackendDetails.copy_exported_program_for_preprocess` lets a backend
decide how that copy is made. The default is `copy.deepcopy`. A backend
can override it to share large constant tensors instead of copying them.
The CUDA backend does that in its low-memory mode, because copying model
weights costs as much memory as the weights themselves.
`EdgeProgramManager.to_backend`, which `to_edge_transform_and_lower`
also uses, already calls the hook. But two other paths do not:
- `to_backend(backend_id, edge_program, compile_specs)` calls
`copy.deepcopy` directly.
- `to_backend(edge_program, partitioner)` calls that same function for
every partition.
On those paths a backend's override is ignored, and every constant is
copied. For a large delegate payload that can be gigabytes of extra
memory during lowering.
## What this change does
The single-backend `to_backend` now calls
`cls.copy_exported_program_for_preprocess(edge_program, compile_specs)`
instead of `copy.deepcopy(edge_program)`.
```python
copied_edge_program = cls.copy_exported_program_for_preprocess(
edge_program, compile_specs
)
```
The default hook is still `copy.deepcopy`, so nothing changes for a
backend that does not override it. Only backends that opted in see a
difference, and they now get the same behavior on every lowering path.
## What was tested
- New unit test in `test_lowered_backend_module.py`: `to_backend` calls
the backend's copy hook once with the program and the compile specs, and
passes what the hook returns to `preprocess`. It fails before this
change (the hook is never called) and passes after it.
- The other tests in that file give the same result before and after.
One end-to-end test fails both ways in my environment, because the
installed runtime does not include the test-only demo backend.1 parent 3ebf896 commit adfd1d7
2 files changed
Lines changed: 24 additions & 1 deletion
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
111 | 111 | | |
112 | 112 | | |
113 | 113 | | |
114 | | - | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
115 | 117 | | |
116 | 118 | | |
117 | 119 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
6 | 6 | | |
7 | 7 | | |
8 | 8 | | |
| 9 | + | |
9 | 10 | | |
10 | 11 | | |
11 | 12 | | |
| |||
279 | 280 | | |
280 | 281 | | |
281 | 282 | | |
| 283 | + | |
| 284 | + | |
| 285 | + | |
| 286 | + | |
| 287 | + | |
| 288 | + | |
| 289 | + | |
| 290 | + | |
| 291 | + | |
| 292 | + | |
| 293 | + | |
| 294 | + | |
| 295 | + | |
| 296 | + | |
| 297 | + | |
| 298 | + | |
| 299 | + | |
| 300 | + | |
| 301 | + | |
| 302 | + | |
0 commit comments