Quantize moveaxis/movedim so they delegate to Ethos-U (#20314)#20314
Quantize moveaxis/movedim so they delegate to Ethos-U (#20314)#20314apullin wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20314
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
|
|
@apullin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D108478011. |
This PR needs a
|
2a0bc8c to
a71d81b
Compare
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Differential Revision: D108478011
JakeStevens
left a comment
There was a problem hiding this comment.
Please add a test for the new annotation.
It would be nice to show that this results in a lowerable chain for a simple model, I believethere are model-level tests in the backend as well
a71d81b to
c822247
Compare
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Differential Revision: D108478011
|
Added a model level test in |
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Differential Revision: D108478011
c822247 to
2658037
Compare
|
let's see what that CI looks like when branch is no longer out of date before landing, the tosa errors are giving me pause so make sure they are resolved |
2658037 to
72945a9
Compare
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Reviewed By: JakeStevens Differential Revision: D108478011
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Reviewed By: JakeStevens Differential Revision: D108478011
72945a9 to
d8e337b
Compare
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Reviewed By: JakeStevens Differential Revision: D108478011
d8e337b to
bd6e94f
Compare
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Reviewed By: JakeStevens Differential Revision: D108478011
bd6e94f to
baf29c1
Compare
|
Something wrong with internal/external sync. No idea. Closing. Will try to regenerate. |
Summary: The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py (_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis therefore leaves those ops unquantized: the quantizer brackets each one with dequantize -> moveaxis(float) -> quantize. On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55 operator-support check (operator_support/ethos_u55_support.py) only delegates permute_copy for int8/int16/int32, so it rejects the float one. Each rejected permute is stranded on the host, splitting the model into many delegated partitions (one NPU island per permute), which bloats the .pte with per-partition delegate overhead and host round-trips. Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec (guarded with getattr for torch-build variance, mirroring the existing transpose.Dimname handling) so they share the input quantization spec exactly like transpose/permute. They then stay int8, decompose to int8 permute_copy, and delegate to the NPU -- eliminating the host float islands. Impact: a quantized example ensemble (ConvNeXt-style blocks that use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers into a single delegate, with zero host permutes and ~24% smaller .pte, with no model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U backend. Reviewed By: JakeStevens Differential Revision: D108478011
Differential Revision: D108478011 Pull Request resolved: #20453
Summary:
The ARM PT2 quantizer's pass-through shared-qspec set in quantization_annotator.py
(_one_to_one_shared_input_qspec) covers permute/permute_copy/transpose/view/squeeze
etc., but omits aten.moveaxis/aten.movedim. A model that uses torch.moveaxis
therefore leaves those ops unquantized: the quantizer brackets each one with
dequantize -> moveaxis(float) -> quantize.
On lowering, moveaxis decomposes to a float permute_copy. The Ethos-U55
operator-support check (operator_support/ethos_u55_support.py) only delegates
permute_copy for int8/int16/int32, so it rejects the float one. Each rejected
permute is stranded on the host, splitting the model into many delegated
partitions (one NPU island per permute), which bloats the .pte with per-partition
delegate overhead and host round-trips.
Add aten.moveaxis.int / aten.movedim.int to _one_to_one_shared_input_qspec
(guarded with getattr for torch-build variance, mirroring the existing
transpose.Dimname handling) so they share the input quantization spec exactly like
transpose/permute. They then stay int8, decompose to int8 permute_copy, and
delegate to the NPU -- eliminating the host float islands.
Impact: a quantized example ensemble (ConvNeXt-style blocks that
use torch.moveaxis) that previously lowered into 9 Ethos-U55 partitions now lowers
into a single delegate, with zero host permutes and ~24% smaller .pte, with no
model changes. Generalizes to any moveaxis/movedim-using model on the Ethos-U
backend.
Reviewed By: JakeStevens
Differential Revision: D108478011