Repository navigation
Add fused_quant, a backend-neutral quantization frontend (#23091) - #23091
mcremon-meta wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23091
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 8321a9b with merge base f2575be ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D119511656. |
This PR needs a
|
Summary: Pull Request resolved: #23091 `fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator that carries its quantization parameters inline. Instead of a quantized graph being a float graph sprinkled with q/dq pairs that each backend has to pattern-match and fold back together, the quantized operator is the unit the graph is built from. Taking a quantized convolution as the example -- before, the quantized op is implicit. Five nodes, and the one doing the arithmetic is float: x_i8 w_i8 bias | | | v v | dequantize dequantize | <- scale, zero_point per edge | | | +-----------+----------+ | v aten.convolution <- float compute | v quantize <- scale, zero_point | v y_i8 After, it is one node, and the qparams ride along with the operands they describe: x_i8 w_i8 bias | | | +-----------+----------+ | v fused_quant.convolution inp : scale, zero_point, dtype, quant_min, quant_max weight : scale, zero_point, dtype, quant_min, quant_max bias : scale, zero_point, dtype, quant_min, quant_max out : scale, zero_point, dtype, quant_min, quant_max | v y_i8 Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form back into the first, so a backend that cannot run a fused op still gets the portable one. This adds that representation under `backends/fused_quant/`. It is ahead-of-time only: a `fused_quant` op has no kernel and is never executed directly. Either a backend claims it and lowers it to its own instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the ordinary portable kernels run it. **Operator representation** (`ops.py`, `ops_utils.py`, `graph_utils.py`) Every `fused_quant` op takes its tensor operands first, then one qparam block (`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per quantizable edge, including the output. A block whose `scale` is `None` means that edge stays in float, so one schema covers the fully quantized, partially quantized and float instances of an op rather than needing a variant for each. Granularity is carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group and blockwise are the same op with differently shaped qparams. **Quantizer** (`quantizer/`) A torchao PT2E `Quantizer` that both annotates and fuses: - `OpQuantizer` annotates one ATen op type and folds matched instances into a `fused_quant` op, given a `QuantizationConfig` and the schema names of its quantized activations, weights and passthrough operands. - `NoopQuantizer` claims nodes with an empty annotation so a later quantizer skips them, leaving the op in float. - `MaxPoolQuantizer` shares the input observer with output 0 of `max_pool2d_with_indices`. - `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a narrower quantizer is listed before the general one it overrides. `preserved_ops()` reports which ATen ops the quantizer matches, so the caller can hold exactly those back from decomposition instead of maintaining a hardcoded list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise specs; `quantizer/defaults.py` provides a default op set, with `get_core_quantizers` for backends that lack per-channel support or the wider op set. **Coloring** (`colorer.py`) `ColorerBase` is the step between optimization and lowering: a backend declares which fused ops it can run (`supported_ops`) and whether a given instance is within its limits (`is_legal`), and the base class tags the rest with `delegation_tag`. The tag is a claim of ownership and the generic passes honour it -- neither `ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps what it claimed, and everything it did not claim falls back to q/dq plus ATen automatically. A backend only has to describe what it supports, not handle what it doesn't. A backend that claims nothing is just the degenerate case of this: every fused op decomposes and the program runs entirely on the ordinary portable ATen kernels. `test_compile.py` has that written out end to end. **Passes** - `pre_quantize_passes/` runs on the traced program before annotation: batch-norm folding, replacing leftover batch norms with convolutions, and recomposing `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel never enters the quantized domain. - `optimization_passes/` runs on the fused edge graph: layout canonicalization, activation hoisting, folding add/mul into linear, quantization absorption, embedding prequantization, linear splitting, requantize formation. - `decompose_fused_quant.py` is the inverse: it turns any fused op back into dequantize + ATen + quantize. It is schema-driven, so it covers every registered fused op rather than a fixed list, handles multi-output ops through `getitem`, and flattens full-rank per-channel qparams to the 1-D form `quantized_decomposed` expects. **Pipeline** (`frontend.py`, `passes.py`, `compile.py`) `compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> passes, and stops. That is the seam a backend plugs into, and it is the only entry point the package offers: from there a backend runs its own colorer, lowering and partitioner, and hands whatever it did not claim to `DecomposeFusedQuant`. The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers that want a different pipeline pass their own, and an empty list skips optimization entirely. One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits the ATen ops that were held back from decomposition, and several of them (`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects them. `compile.py` feeds `quantizer.preserved_ops()` into `EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by construction: the ops held back are precisely the ops re-emitted. **Known gaps**, also in the README: the package depends on `torchao`, which is not yet a core ExecuTorch dependency; and the optimization pass group still imports about fifteen passes from `backends/cadence/aot` that are not Cadence-specific and would be better placed in `backends/transforms`. Differential Revision: D119511656
160be79 to
82f494a
Compare
Summary: Pull Request resolved: #23091 `fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator that carries its quantization parameters inline. Instead of a quantized graph being a float graph sprinkled with q/dq pairs that each backend has to pattern-match and fold back together, the quantized operator is the unit the graph is built from. Taking a quantized convolution as the example -- before, the quantized op is implicit. Five nodes, and the one doing the arithmetic is float: x_i8 w_i8 bias | | | v v | dequantize dequantize | <- scale, zero_point per edge | | | +-----------+----------+ | v aten.convolution <- float compute | v quantize <- scale, zero_point | v y_i8 After, it is one node, and the qparams ride along with the operands they describe: x_i8 w_i8 bias | | | +-----------+----------+ | v fused_quant.convolution inp : scale, zero_point, dtype, quant_min, quant_max weight : scale, zero_point, dtype, quant_min, quant_max bias : scale, zero_point, dtype, quant_min, quant_max out : scale, zero_point, dtype, quant_min, quant_max | v y_i8 Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form back into the first, so a backend that cannot run a fused op still gets the portable one. This adds that representation under `backends/fused_quant/`. It is ahead-of-time only: a `fused_quant` op has no kernel and is never executed directly. Either a backend claims it and lowers it to its own instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the ordinary portable kernels run it. **Operator representation** (`ops.py`, `ops_utils.py`, `graph_utils.py`) Every `fused_quant` op takes its tensor operands first, then one qparam block (`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per quantizable edge, including the output. A block whose `scale` is `None` means that edge stays in float, so one schema covers the fully quantized, partially quantized and float instances of an op rather than needing a variant for each. Granularity is carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group and blockwise are the same op with differently shaped qparams. **Quantizer** (`quantizer/`) A torchao PT2E `Quantizer` that both annotates and fuses: - `OpQuantizer` annotates one ATen op type and folds matched instances into a `fused_quant` op, given a `QuantizationConfig` and the schema names of its quantized activations, weights and passthrough operands. - `NoopQuantizer` claims nodes with an empty annotation so a later quantizer skips them, leaving the op in float. - `MaxPoolQuantizer` shares the input observer with output 0 of `max_pool2d_with_indices`. - `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a narrower quantizer is listed before the general one it overrides. `preserved_ops()` reports which ATen ops the quantizer matches, so the caller can hold exactly those back from decomposition instead of maintaining a hardcoded list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise specs; `quantizer/defaults.py` provides a default op set, with `get_core_quantizers` for backends that lack per-channel support or the wider op set. **Coloring** (`colorer.py`) `ColorerBase` is the step between optimization and lowering: a backend declares which fused ops it can run (`supported_ops`) and whether a given instance is within its limits (`is_legal`), and the base class tags the rest with `delegation_tag`. The tag is a claim of ownership and the generic passes honour it -- neither `ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps what it claimed, and everything it did not claim falls back to q/dq plus ATen automatically. A backend only has to describe what it supports, not handle what it doesn't. A backend that claims nothing is just the degenerate case of this: every fused op decomposes and the program runs entirely on the ordinary portable ATen kernels. `test_compile.py` has that written out end to end. **Passes** - `pre_quantize_passes/` runs on the traced program before annotation: batch-norm folding, replacing leftover batch norms with convolutions, and recomposing `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel never enters the quantized domain. - `optimization_passes/` runs on the fused edge graph: layout canonicalization, activation hoisting, folding add/mul into linear, quantization absorption, embedding prequantization, linear splitting, requantize formation. - `decompose_fused_quant.py` is the inverse: it turns any fused op back into dequantize + ATen + quantize. It is schema-driven, so it covers every registered fused op rather than a fixed list, handles multi-output ops through `getitem`, and flattens full-rank per-channel qparams to the 1-D form `quantized_decomposed` expects. **Pipeline** (`frontend.py`, `passes.py`, `compile.py`) `compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> passes, and stops. That is the seam a backend plugs into, and it is the only entry point the package offers: from there a backend runs its own colorer, lowering and partitioner, and hands whatever it did not claim to `DecomposeFusedQuant`. The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers that want a different pipeline pass their own, and an empty list skips optimization entirely. One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits the ATen ops that were held back from decomposition, and several of them (`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects them. `compile.py` feeds `quantizer.preserved_ops()` into `EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by construction: the ops held back are precisely the ops re-emitted. **Known gaps**, also in the README: the package depends on `torchao`, which is not yet a core ExecuTorch dependency; and the optimization pass group still imports about fifteen passes from `backends/cadence/aot` that are not Cadence-specific and would be better placed in `backends/transforms`. Differential Revision: D119511656
82f494a to
f16d602
Compare
| ## Directory structure | ||
|
|
||
| ``` | ||
| backends/fused_quant/ |
There was a problem hiding this comment.
| backends/fused_quant/ | |
| backends/transforms/fused_quant/ |
| @@ -0,0 +1,68 @@ | |||
| # fused_quant | |||
|
|
|||
| `fused_quant` is a backend-neutral quantization frontend for ExecuTorch. It represents a | |||
There was a problem hiding this comment.
What are your thoughts on making this extensible so backends can customize to their needs? Because that's what they do today to match their hardware capabilities, hence different quantizers etc. There might be overlap which this can help dedup, but I do expect some need for specialization per backend if this were to adopted widely, IMHO.
There was a problem hiding this comment.
of course, that's what all of the "adoption" follow ups will do. This is just landing the infra layers
| @@ -0,0 +1,68 @@ | |||
| # fused_quant | |||
|
|
|||
| `fused_quant` is a backend-neutral quantization frontend for ExecuTorch. It represents a | |||
There was a problem hiding this comment.
Nit: rather than calling this fused_quant, this should reflect the quant_scheme(s) it reflects or something else like "some_name_standard_int8_symmetric_quant_IR" (obv not that) but something like this?
There was a problem hiding this comment.
well the whole point of this layer is that it supports all of those schemes natively, together. Including fp32 actually, where you would "just" get the optimizations passes and layout changes. cc @JakeStevens since we talked about that briefly
There was a problem hiding this comment.
Maybe a bit of bike shedding but the fused just leads my thoughts to fusing functional ops. Maybe just quantized_ir ?
| return val.fake_mode | ||
|
|
||
|
|
||
| def add_constant( |
There was a problem hiding this comment.
Probably should reuse constant creation functions in backends/transforms/utils.py
There was a problem hiding this comment.
Probably should reuse constant creation functions in backends/transforms/utils.py
yes, let me see if we have more of those opportunities as well
| @@ -0,0 +1,25 @@ | |||
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |||
There was a problem hiding this comment.
why not put this in exir.passes?
There was a problem hiding this comment.
we can do that for some of those for sure. Let me see if there are more candidates
|
|
||
| The rewrites run here rather than in the caller because they must happen | ||
| after tracing (they match on ATen ops) but before annotation (the quantizer | ||
| must see the rewritten ops). Which ones run is derived from the quantizer: |
| fold_quantize_into_mutable_buffers=fold_quantize_into_mutable_buffers, | ||
| ) | ||
| _normalize_affine_overloads(gm) | ||
| # convert_pt2e rewrites the graph without refreshing meta["val"], so shape and |
There was a problem hiding this comment.
perhaps we should fix it in the convert?
|
|
||
|
|
||
| @torch.no_grad() | ||
| def trace( |
|
@claude review this PR. Focus on architectural aspects of the design and how this is setup to be adopted by existing backends as an IR Toolkit for efficient lowering, while allowing them to be opinionated about their hardware. Also other standard ET features like Quantizer composability and same for partitioner - if someone were to write a recipe for a composed setup. Lastly custom op versioning, BC contract, QAT support, for the dialect. How can optimization passes can be split and opt-in at a finer granularity? Also does exir dir makes a better home for this? |
|
Claude finished @digantdesai's task in 1m 38s —— View job Summary: The fused-op IR idea is sound, but this PR isn't ready to serve as a shared toolkit for backends yet. It only works through its own Detailed reviewPR Review: #23091SummaryThis adds API DesignOp dialect, versioning and BC
Adopting it from a backend (colorer → partitioner)
Quantizer composability
Recipe /
QAT
Pass granularity (how to make passes opt-in)
Infrastructure
Code Quality
Testing
RecommendationRequest Changes Blocking:
Needs discussion before backends adopt this:
|
f16d602 to
1f5540f
Compare
Summary: Pull Request resolved: #23091 `fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator that carries its quantization parameters inline. Instead of a quantized graph being a float graph sprinkled with q/dq pairs that each backend has to pattern-match and fold back together, the quantized operator is the unit the graph is built from. Taking a quantized convolution as the example -- before, the quantized op is implicit. Five nodes, and the one doing the arithmetic is float: x_i8 w_i8 bias | | | v v | dequantize dequantize | <- scale, zero_point per edge | | | +-----------+----------+ | v aten.convolution <- float compute | v quantize <- scale, zero_point | v y_i8 After, it is one node, and the qparams ride along with the operands they describe: x_i8 w_i8 bias | | | +-----------+----------+ | v fused_quant.convolution inp : scale, zero_point, dtype, quant_min, quant_max weight : scale, zero_point, dtype, quant_min, quant_max bias : scale, zero_point, dtype, quant_min, quant_max out : scale, zero_point, dtype, quant_min, quant_max | v y_i8 Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form back into the first, so a backend that cannot run a fused op still gets the portable one. `test_op_numerics` checks this exactly, for every registered op. This adds that representation under `backends/fused_quant/`. It is ahead-of-time only: each op has a reference (eager) kernel that defines its numerics, but no runtime kernel, so a `fused_quant` op never reaches `to_executorch`. Either a backend claims it and lowers it to its own instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails the compile if one leaks. **Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`) Every `fused_quant` op takes its tensor operands first, then one qparam block (`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand, then the `out` block, then the ATen op's remaining arguments in ATen order. A block whose `scale` is `None` means that edge stays in float, so one schema covers the fully quantized, partially quantized and float instances of an op rather than needing a variant for each. For an input block `dtype` is the dequantized float dtype; for `out` it is the quantized storage dtype. Granularity is carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group and blockwise are the same op with differently shaped qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` / `is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel. The README writes down what is stable for backends: the argument layout and field meanings are frozen, a new capability is a new overload or a new trailing argument with a behavior-preserving default, and the ops are never serialized. Fake kernels report the dtype the eager kernel actually returns, including for a partially quantized op (quantized input, float output), where they previously reported the int8 storage dtype. **Quantizer** (`quantizer/`) A torchao PT2E `Quantizer` that both annotates and fuses: - `OpQuantizer` annotates one ATen op type and folds matched instances into a `fused_quant` op, given a `QuantizationConfig` and the schema names of its quantized activations, weights and passthrough operands. - `NoopQuantizer` claims nodes with an empty annotation so a later quantizer skips them, leaving the op in float. - `MaxPoolQuantizer` shares the input observer with output 0 of `max_pool2d_with_indices`. - `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a narrower quantizer is listed before the general one it overrides. Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK, Arm) check before claiming a node, and `fuse()` only folds nodes the same quantizer annotated -- never one claimed by another quantizer. `preserved_ops()` reports which ATen ops the quantizer matches, so the caller can hold exactly those back from decomposition instead of maintaining a hardcoded list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static specs (dynamic qspecs are not offered, since fusion needs constant qparams); `quantizer/defaults.py` provides a default op set, with `get_core_quantizers` for backends that lack per-channel support or the wider op set. Weights default to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A transposed convolution gets its per-channel weight scales on axis 1, where its output channels are, or per-tensor weights when `groups > 1`. **Coloring** (`colorer.py`) `ColorerBase` is the step between optimization and lowering: a backend declares which fused ops it can run (`supported_ops`) and whether a given instance is within its limits (`is_legal`), and the base class tags the rest with `delegation_tag`. Colorers compose by priority: the first to claim a node keeps it, and a later colorer skips already-tagged nodes. The tag is a claim of ownership and the generic passes honour it -- neither `ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps what it claimed, and everything it did not claim falls back to q/dq plus ATen automatically. A colored node the partitioner then rejects is caught by running `DecomposeFusedQuant` again after `to_backend`. A backend that claims nothing is just the degenerate case of this: every fused op decomposes and the program runs entirely on the ordinary portable ATen kernels. `test_compile.py` has that written out end to end. **Passes** - `pre_quantize_passes/` runs on the traced program before annotation: batch-norm folding, replacing leftover batch norms with convolutions, and recomposing `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel never enters the quantized domain. - `optimization_passes/` runs on the fused edge graph: layout canonicalization, activation hoisting, folding add/mul into linear, quantization absorption, embedding prequantization, linear splitting, requantize formation. `IterativePassGroup` logs a warning when a group hits its step cap. - `decompose_fused_quant.py` is the inverse: it turns any fused op back into dequantize + ATen + quantize. It is schema-driven, so it covers every registered fused op rather than a fixed list, and handles multi-output ops through `getitem`. It validates every node before rewriting any, so an unsupported one (grouped/blockwise qparams) fails the pass without leaving a partly decomposed graph. Per-channel qparams are emitted as 1-D constants rather than a `view_copy` of the full-rank one, and the qparam constants left dead are removed. **Shared passes moved to `backends/transforms`** The optimization group no longer depends on any backend. It used 22 passes from `backends/cadence/aot`; none were Cadence-specific in what they do: - Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`, `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a `backends/transforms` pass, which is now imported directly. - The other 19 move to `backends/transforms`, one module and target each. The ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now take those op sets as class attributes holding only the standard ops. Cadence keeps each pass under its old name as a thin subclass that adds its own ops back, so Cadence's behavior is unchanged. - `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose qparams differ, and there is no backend-neutral one. The shared pass takes a `requantize_op` (default `None`) and defaults to `allow_requantize=False`; Cadence sets its op and keeps `True`. - `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and `quantize_reorder_utils` holds the two op sets the reorder passes share. Tests no longer import Cadence either: `test/helpers.trace` replaces `cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone. `buck2 uquery` finds no `backends/cadence` target in the dependency closure of any `fused_quant` target. **Pipeline** (`frontend.py`, `passes.py`, `compile.py`) `compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> passes, and stops. That is the seam a backend plugs into, and it is the only entry point the package offers: from there a backend runs its own colorer, lowering and partitioner, and hands whatever it did not claim to `DecomposeFusedQuant`. `trace(..., is_qat=True)` exports in training mode and keeps batch norms, for `prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded. Either way the caller's model is left in the mode it was passed in. The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers that want a different pipeline pass their own, and an empty list skips optimization entirely. One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits the ATen ops that were held back from decomposition, and several of them (`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects them. `compile.py` feeds `quantizer.preserved_ops()` into `EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by construction: the ops held back are precisely the ops re-emitted. **Known gaps** - The package depends on `torchao`, which is not yet a core ExecuTorch dependency. - `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting `dequantize_affine`, so an unclaimed blockwise op fails the compile. - Bias stays float in the default configs; int32 bias is neither configurable nor documented yet. - Op registration happens on import, and `passes.py` looks up `exir_ops.edge.fused_quant.*` at module scope. - Follow-ups from review, deliberately not in this diff: splitting the pass pipeline into opt-in groups (which takes `ToChannelsLast` out of the default), a colored-node `Partitioner` adapter, fitting the quantizer and frontend into `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating schemas from the ATen ones, and the core/transforms/quantizer placement split. Differential Revision: D119511656
1f5540f to
3e4896f
Compare
Summary: Pull Request resolved: #23091 `fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator that carries its quantization parameters inline. Instead of a quantized graph being a float graph sprinkled with q/dq pairs that each backend has to pattern-match and fold back together, the quantized operator is the unit the graph is built from. Taking a quantized convolution as the example -- before, the quantized op is implicit. Five nodes, and the one doing the arithmetic is float: x_i8 w_i8 bias | | | v v | dequantize dequantize | <- scale, zero_point per edge | | | +-----------+----------+ | v aten.convolution <- float compute | v quantize <- scale, zero_point | v y_i8 After, it is one node, and the qparams ride along with the operands they describe: x_i8 w_i8 bias | | | +-----------+----------+ | v fused_quant.convolution inp : scale, zero_point, dtype, quant_min, quant_max weight : scale, zero_point, dtype, quant_min, quant_max bias : scale, zero_point, dtype, quant_min, quant_max out : scale, zero_point, dtype, quant_min, quant_max | v y_i8 Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form back into the first, so a backend that cannot run a fused op still gets the portable one. `test_op_numerics` checks this exactly, for every registered op. This adds that representation under `backends/fused_quant/`. It is ahead-of-time only: each op has a reference (eager) kernel that defines its numerics, but no runtime kernel, so a `fused_quant` op never reaches `to_executorch`. Either a backend claims it and lowers it to its own instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails the compile if one leaks. **Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`) Every `fused_quant` op takes its tensor operands first, then one qparam block (`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand, then the `out` block, then the ATen op's remaining arguments in ATen order. A block whose `scale` is `None` means that edge stays in float, so one schema covers the fully quantized, partially quantized and float instances of an op rather than needing a variant for each. For an input block `dtype` is the dequantized float dtype; for `out` it is the quantized storage dtype. Granularity is carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group and blockwise are the same op with differently shaped qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` / `is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel. The README writes down what is stable for backends: the argument layout and field meanings are frozen, a new capability is a new overload or a new trailing argument with a behavior-preserving default, and the ops are never serialized. Fake kernels report the dtype the eager kernel actually returns, including for a partially quantized op (quantized input, float output), where they previously reported the int8 storage dtype. **Quantizer** (`quantizer/`) A torchao PT2E `Quantizer` that both annotates and fuses: - `OpQuantizer` annotates one ATen op type and folds matched instances into a `fused_quant` op, given a `QuantizationConfig` and the schema names of its quantized activations, weights and passthrough operands. - `NoopQuantizer` claims nodes with an empty annotation so a later quantizer skips them, leaving the op in float. - `MaxPoolQuantizer` shares the input observer with output 0 of `max_pool2d_with_indices`. - `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a narrower quantizer is listed before the general one it overrides. Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK, Arm) check before claiming a node, and `fuse()` only folds nodes the same quantizer annotated -- never one claimed by another quantizer. `preserved_ops()` reports which ATen ops the quantizer matches, so the caller can hold exactly those back from decomposition instead of maintaining a hardcoded list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static specs (dynamic qspecs are not offered, since fusion needs constant qparams); `quantizer/defaults.py` provides a default op set, with `get_core_quantizers` for backends that lack per-channel support or the wider op set. Weights default to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A transposed convolution gets its per-channel weight scales on axis 1, where its output channels are, or per-tensor weights when `groups > 1`. **Coloring** (`colorer.py`) `ColorerBase` is the step between optimization and lowering: a backend declares which fused ops it can run (`supported_ops`) and whether a given instance is within its limits (`is_legal`), and the base class tags the rest with `delegation_tag`. Colorers compose by priority: the first to claim a node keeps it, and a later colorer skips already-tagged nodes. The tag is a claim of ownership and the generic passes honour it -- neither `ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps what it claimed, and everything it did not claim falls back to q/dq plus ATen automatically. A colored node the partitioner then rejects is caught by running `DecomposeFusedQuant` again after `to_backend`. A backend that claims nothing is just the degenerate case of this: every fused op decomposes and the program runs entirely on the ordinary portable ATen kernels. `test_compile.py` has that written out end to end. **Passes** - `pre_quantize_passes/` runs on the traced program before annotation: batch-norm folding, replacing leftover batch norms with convolutions, and recomposing `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel never enters the quantized domain. - `optimization_passes/` runs on the fused edge graph: layout canonicalization, activation hoisting, folding add/mul into linear, quantization absorption, embedding prequantization, linear splitting, requantize formation. `IterativePassGroup` logs a warning when a group hits its step cap. - `decompose_fused_quant.py` is the inverse: it turns any fused op back into dequantize + ATen + quantize. It is schema-driven, so it covers every registered fused op rather than a fixed list, and handles multi-output ops through `getitem`. It validates every node before rewriting any, so an unsupported one (grouped/blockwise qparams) fails the pass without leaving a partly decomposed graph. Per-channel qparams are emitted as 1-D constants rather than a `view_copy` of the full-rank one, and the qparam constants left dead are removed. **Shared passes moved to `backends/transforms`** The optimization group no longer depends on any backend. It used 22 passes from `backends/cadence/aot`; none were Cadence-specific in what they do: - Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`, `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a `backends/transforms` pass, which is now imported directly. - The other 19 move to `backends/transforms`, one module and target each. The ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now take those op sets as class attributes holding only the standard ops. Cadence keeps each pass under its old name as a thin subclass that adds its own ops back, so Cadence's behavior is unchanged. - `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose qparams differ, and there is no backend-neutral one. The shared pass takes a `requantize_op` (default `None`) and defaults to `allow_requantize=False`; Cadence sets its op and keeps `True`. - `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and `quantize_reorder_utils` holds the two op sets the reorder passes share. Tests no longer import Cadence either: `test/helpers.trace` replaces `cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone. `buck2 uquery` finds no `backends/cadence` target in the dependency closure of any `fused_quant` target. **Pipeline** (`frontend.py`, `passes.py`, `compile.py`) `compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> passes, and stops. That is the seam a backend plugs into, and it is the only entry point the package offers: from there a backend runs its own colorer, lowering and partitioner, and hands whatever it did not claim to `DecomposeFusedQuant`. `trace(..., is_qat=True)` exports in training mode and keeps batch norms, for `prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded. Either way the caller's model is left in the mode it was passed in. The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers that want a different pipeline pass their own, and an empty list skips optimization entirely. One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits the ATen ops that were held back from decomposition, and several of them (`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects them. `compile.py` feeds `quantizer.preserved_ops()` into `EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by construction: the ops held back are precisely the ops re-emitted. **Known gaps** - The package depends on `torchao`, which is not yet a core ExecuTorch dependency. - `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting `dequantize_affine`, so an unclaimed blockwise op fails the compile. - Bias stays float in the default configs; int32 bias is neither configurable nor documented yet. - Op registration happens on import, and `passes.py` looks up `exir_ops.edge.fused_quant.*` at module scope. - Follow-ups from review, deliberately not in this diff: splitting the pass pipeline into opt-in groups (which takes `ToChannelsLast` out of the default), a colored-node `Partitioner` adapter, fitting the quantizer and frontend into `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating schemas from the ATen ones, and the core/transforms/quantizer placement split. Differential Revision: D119511656
3e4896f to
4b28cb0
Compare
Summary: Pull Request resolved: #23091 `fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator that carries its quantization parameters inline. Instead of a quantized graph being a float graph sprinkled with q/dq pairs that each backend has to pattern-match and fold back together, the quantized operator is the unit the graph is built from. Taking a quantized convolution as the example -- before, the quantized op is implicit. Five nodes, and the one doing the arithmetic is float: x_i8 w_i8 bias | | | v v | dequantize dequantize | <- scale, zero_point per edge | | | +-----------+----------+ | v aten.convolution <- float compute | v quantize <- scale, zero_point | v y_i8 After, it is one node, and the qparams ride along with the operands they describe: x_i8 w_i8 bias | | | +-----------+----------+ | v fused_quant.convolution inp : scale, zero_point, dtype, quant_min, quant_max weight : scale, zero_point, dtype, quant_min, quant_max bias : scale, zero_point, dtype, quant_min, quant_max out : scale, zero_point, dtype, quant_min, quant_max | v y_i8 Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form back into the first, so a backend that cannot run a fused op still gets the portable one. `test_op_numerics` checks this exactly, for every registered op. This adds that representation under `backends/fused_quant/`. It is ahead-of-time only: each op has a reference (eager) kernel that defines its numerics, but no runtime kernel, so a `fused_quant` op never reaches `to_executorch`. Either a backend claims it and lowers it to its own instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails the compile if one leaks. **Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`) Every `fused_quant` op takes its tensor operands first, then one qparam block (`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand, then the `out` block, then the ATen op's remaining arguments in ATen order. A block whose `scale` is `None` means that edge stays in float, so one schema covers the fully quantized, partially quantized and float instances of an op rather than needing a variant for each. For an input block `dtype` is the dequantized float dtype; for `out` it is the quantized storage dtype. Granularity is carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group and blockwise are the same op with differently shaped qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` / `is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel. The README writes down what is stable for backends: the argument layout and field meanings are frozen, a new capability is a new overload or a new trailing argument with a behavior-preserving default, and the ops are never serialized. Fake kernels report the dtype the eager kernel actually returns, including for a partially quantized op (quantized input, float output), where they previously reported the int8 storage dtype. **Quantizer** (`quantizer/`) A torchao PT2E `Quantizer` that both annotates and fuses: - `OpQuantizer` annotates one ATen op type and folds matched instances into a `fused_quant` op, given a `QuantizationConfig` and the schema names of its quantized activations, weights and passthrough operands. - `NoopQuantizer` claims nodes with an empty annotation so a later quantizer skips them, leaving the op in float. - `MaxPoolQuantizer` shares the input observer with output 0 of `max_pool2d_with_indices`. - `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a narrower quantizer is listed before the general one it overrides. Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK, Arm) check before claiming a node, and `fuse()` only folds nodes the same quantizer annotated -- never one claimed by another quantizer. `preserved_ops()` reports which ATen ops the quantizer matches, so the caller can hold exactly those back from decomposition instead of maintaining a hardcoded list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static specs (dynamic qspecs are not offered, since fusion needs constant qparams); `quantizer/defaults.py` provides a default op set, with `get_core_quantizers` for backends that lack per-channel support or the wider op set. Weights default to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A transposed convolution gets its per-channel weight scales on axis 1, where its output channels are, or per-tensor weights when `groups > 1`. **Coloring** (`colorer.py`) `ColorerBase` is the step between optimization and lowering: a backend declares which fused ops it can run (`supported_ops`) and whether a given instance is within its limits (`is_legal`), and the base class tags the rest with `delegation_tag`. Colorers compose by priority: the first to claim a node keeps it, and a later colorer skips already-tagged nodes. The tag is a claim of ownership and the generic passes honour it -- neither `ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps what it claimed, and everything it did not claim falls back to q/dq plus ATen automatically. A colored node the partitioner then rejects is caught by running `DecomposeFusedQuant` again after `to_backend`. A backend that claims nothing is just the degenerate case of this: every fused op decomposes and the program runs entirely on the ordinary portable ATen kernels. `test_compile.py` has that written out end to end. **Passes** - `pre_quantize_passes/` runs on the traced program before annotation: batch-norm folding, replacing leftover batch norms with convolutions, and recomposing `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel never enters the quantized domain. - `optimization_passes/` runs on the fused edge graph: layout canonicalization, activation hoisting, folding add/mul into linear, quantization absorption, embedding prequantization, linear splitting, requantize formation. `IterativePassGroup` logs a warning when a group hits its step cap. - `decompose_fused_quant.py` is the inverse: it turns any fused op back into dequantize + ATen + quantize. It is schema-driven, so it covers every registered fused op rather than a fixed list, and handles multi-output ops through `getitem`. It validates every node before rewriting any, so an unsupported one (grouped/blockwise qparams) fails the pass without leaving a partly decomposed graph. Per-channel qparams are emitted as 1-D constants rather than a `view_copy` of the full-rank one, and the qparam constants left dead are removed. **Shared passes moved to `backends/transforms`** The optimization group no longer depends on any backend. It used 22 passes from `backends/cadence/aot`; none were Cadence-specific in what they do: - Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`, `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a `backends/transforms` pass, which is now imported directly. - The other 19 move to `backends/transforms`, one module and target each. The ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now take those op sets as class attributes holding only the standard ops. Cadence keeps each pass under its old name as a thin subclass that adds its own ops back, so Cadence's behavior is unchanged. - `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose qparams differ, and there is no backend-neutral one. The shared pass takes a `requantize_op` (default `None`) and defaults to `allow_requantize=False`; Cadence sets its op and keeps `True`. - `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and `quantize_reorder_utils` holds the two op sets the reorder passes share. Tests no longer import Cadence either: `test/helpers.trace` replaces `cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone. `buck2 uquery` finds no `backends/cadence` target in the dependency closure of any `fused_quant` target. **Pipeline** (`frontend.py`, `passes.py`, `compile.py`) `compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> passes, and stops. That is the seam a backend plugs into, and it is the only entry point the package offers: from there a backend runs its own colorer, lowering and partitioner, and hands whatever it did not claim to `DecomposeFusedQuant`. `trace(..., is_qat=True)` exports in training mode and keeps batch norms, for `prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded. Either way the caller's model is left in the mode it was passed in. The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers that want a different pipeline pass their own, and an empty list skips optimization entirely. One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits the ATen ops that were held back from decomposition, and several of them (`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects them. `compile.py` feeds `quantizer.preserved_ops()` into `EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by construction: the ops held back are precisely the ops re-emitted. **Known gaps** - The package depends on `torchao`, which is not yet a core ExecuTorch dependency. - `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting `dequantize_affine`, so an unclaimed blockwise op fails the compile. - Bias stays float in the default configs; int32 bias is neither configurable nor documented yet. - Op registration happens on import, and `passes.py` looks up `exir_ops.edge.fused_quant.*` at module scope. - Follow-ups from review, deliberately not in this diff: splitting the pass pipeline into opt-in groups (which takes `ToChannelsLast` out of the default), a colored-node `Partitioner` adapter, fitting the quantizer and frontend into `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating schemas from the ATen ones, and the core/transforms/quantizer placement split. Differential Revision: D119511656
4b28cb0 to
3ee0b0c
Compare
Summary: Pull Request resolved: #23091 `fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single `torch.ops.fused_quant.*` operator that carries its quantization parameters inline. Instead of a quantized graph being a float graph sprinkled with q/dq pairs that each backend has to pattern-match and fold back together, the quantized operator is the unit the graph is built from. Taking a quantized convolution as the example -- before, the quantized op is implicit. Five nodes, and the one doing the arithmetic is float: x_i8 w_i8 bias | | | v v | dequantize dequantize | <- scale, zero_point per edge | | | +-----------+----------+ | v aten.convolution <- float compute | v quantize <- scale, zero_point | v y_i8 After, it is one node, and the qparams ride along with the operands they describe: x_i8 w_i8 bias | | | +-----------+----------+ | v fused_quant.convolution inp : scale, zero_point, dtype, quant_min, quant_max weight : scale, zero_point, dtype, quant_min, quant_max bias : scale, zero_point, dtype, quant_min, quant_max out : scale, zero_point, dtype, quant_min, quant_max | v y_i8 Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form back into the first, so a backend that cannot run a fused op still gets the portable one. `test_op_numerics` checks this exactly, for every registered op. This adds that representation under `backends/fused_quant/`. It is ahead-of-time only: each op has a reference (eager) kernel that defines its numerics, but no runtime kernel, so a `fused_quant` op never reaches `to_executorch`. Either a backend claims it and lowers it to its own instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails the compile if one leaks. **Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`) Every `fused_quant` op takes its tensor operands first, then one qparam block (`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand, then the `out` block, then the ATen op's remaining arguments in ATen order. A block whose `scale` is `None` means that edge stays in float, so one schema covers the fully quantized, partially quantized and float instances of an op rather than needing a variant for each. For an input block `dtype` is the dequantized float dtype; for `out` it is the quantized storage dtype. Granularity is carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group and blockwise are the same op with differently shaped qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` / `is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel. The README writes down what is stable for backends: the argument layout and field meanings are frozen, a new capability is a new overload or a new trailing argument with a behavior-preserving default, and the ops are never serialized. Fake kernels report the dtype the eager kernel actually returns, including for a partially quantized op (quantized input, float output), where they previously reported the int8 storage dtype. **Quantizer** (`quantizer/`) A torchao PT2E `Quantizer` that both annotates and fuses: - `OpQuantizer` annotates one ATen op type and folds matched instances into a `fused_quant` op, given a `QuantizationConfig` and the schema names of its quantized activations, weights and passthrough operands. - `NoopQuantizer` claims nodes with an empty annotation so a later quantizer skips them, leaving the op in float. - `MaxPoolQuantizer` shares the input observer with output 0 of `max_pool2d_with_indices`. - `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a narrower quantizer is listed before the general one it overrides. Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK, Arm) check before claiming a node, and `fuse()` only folds nodes the same quantizer annotated -- never one claimed by another quantizer. `preserved_ops()` reports which ATen ops the quantizer matches, so the caller can hold exactly those back from decomposition instead of maintaining a hardcoded list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static specs (dynamic qspecs are not offered, since fusion needs constant qparams); `quantizer/defaults.py` provides a default op set, with `get_core_quantizers` for backends that lack per-channel support or the wider op set. Weights default to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A transposed convolution gets its per-channel weight scales on axis 1, where its output channels are, or per-tensor weights when `groups > 1`. **Coloring** (`colorer.py`) `ColorerBase` is the step between optimization and lowering: a backend declares which fused ops it can run (`supported_ops`) and whether a given instance is within its limits (`is_legal`), and the base class tags the rest with `delegation_tag`. Colorers compose by priority: the first to claim a node keeps it, and a later colorer skips already-tagged nodes. The tag is a claim of ownership and the generic passes honour it -- neither `ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps what it claimed, and everything it did not claim falls back to q/dq plus ATen automatically. A colored node the partitioner then rejects is caught by running `DecomposeFusedQuant` again after `to_backend`. A backend that claims nothing is just the degenerate case of this: every fused op decomposes and the program runs entirely on the ordinary portable ATen kernels. `test_compile.py` has that written out end to end. **Passes** - `pre_quantize_passes/` runs on the traced program before annotation: batch-norm folding, replacing leftover batch norms with convolutions, and recomposing `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel never enters the quantized domain. - `optimization_passes/` runs on the fused edge graph: layout canonicalization, activation hoisting, folding add/mul into linear, quantization absorption, embedding prequantization, linear splitting, requantize formation. `IterativePassGroup` logs a warning when a group hits its step cap. - `decompose_fused_quant.py` is the inverse: it turns any fused op back into dequantize + ATen + quantize. It is schema-driven, so it covers every registered fused op rather than a fixed list, and handles multi-output ops through `getitem`. It validates every node before rewriting any, so an unsupported one (grouped/blockwise qparams) fails the pass without leaving a partly decomposed graph. Per-channel qparams are emitted as 1-D constants rather than a `view_copy` of the full-rank one, and the qparam constants left dead are removed. **Shared passes moved to `backends/transforms`** The optimization group no longer depends on any backend. It used 22 passes from `backends/cadence/aot`; none were Cadence-specific in what they do: - Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`, `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a `backends/transforms` pass, which is now imported directly. - The other 19 move to `backends/transforms`, one module and target each. The ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now take those op sets as class attributes holding only the standard ops. Cadence keeps each pass under its old name as a thin subclass that adds its own ops back, so Cadence's behavior is unchanged. - `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose qparams differ, and there is no backend-neutral one. The shared pass takes a `requantize_op` (default `None`) and defaults to `allow_requantize=False`; Cadence sets its op and keeps `True`. - `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and `quantize_reorder_utils` holds the two op sets the reorder passes share. Tests no longer import Cadence either: `test/helpers.trace` replaces `cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone. `buck2 uquery` finds no `backends/cadence` target in the dependency closure of any `fused_quant` target. **Pipeline** (`frontend.py`, `passes.py`, `compile.py`) `compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse -> export -> to_edge -> passes, and stops. That is the seam a backend plugs into, and it is the only entry point the package offers: from there a backend runs its own colorer, lowering and partitioner, and hands whatever it did not claim to `DecomposeFusedQuant`. `trace(..., is_qat=True)` exports in training mode and keeps batch norms, for `prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded. Either way the caller's model is left in the mode it was passed in. The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers that want a different pipeline pass their own, and an empty list skips optimization entirely. One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits the ATen ops that were held back from decomposition, and several of them (`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects them. `compile.py` feeds `quantizer.preserved_ops()` into `EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by construction: the ops held back are precisely the ops re-emitted. **Known gaps** - The package depends on `torchao`, which is not yet a core ExecuTorch dependency. - `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting `dequantize_affine`, so an unclaimed blockwise op fails the compile. - Bias stays float in the default configs; int32 bias is neither configurable nor documented yet. - Op registration happens on import, and `passes.py` looks up `exir_ops.edge.fused_quant.*` at module scope. - Follow-ups from review, deliberately not in this diff: splitting the pass pipeline into opt-in groups (which takes `ToChannelsLast` out of the default), a colored-node `Partitioner` adapter, fitting the quantizer and frontend into `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating schemas from the ATen ones, and the core/transforms/quantizer placement split. Differential Revision: D119511656
3ee0b0c to
8321a9b
Compare
Summary:
fused_quantrepresents adequantize -> aten op -> quantizetriple as a singletorch.ops.fused_quant.*operator that carries its quantization parametersinline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.
Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:
After, it is one node, and the qparams ride along with the operands they
describe:
Nothing is lost in the round trip:
DecomposeFusedQuantturns the second formback into the first, so a backend that cannot run a fused op still gets the
portable one.
test_op_numericschecks this exactly, for every registered op.This adds that representation under
backends/fused_quant/. It isahead-of-time only: each op has a reference (eager) kernel that defines its
numerics, but no runtime kernel, so a
fused_quantop never reachesto_executorch. Either a backend claims it and lowers it to its owninstruction, or
DecomposeFusedQuantturns it back into the q/dq + ATen formabove and the ordinary portable kernels run it.
AssertNoFusedQuantOpsfailsthe compile if one leaks.
Operator representation and contract (
ops.py,ops_utils.py,graph_utils.py)Every
fused_quantop takes its tensor operands first, then one qparam block(
scale,zero_point,dtype,quant_min,quant_max) per tensor operand,then the
outblock, then the ATen op's remaining arguments in ATen order. Ablock whose
scaleisNonemeans that edge stays in float, so one schemacovers the fully quantized, partially quantized and float instances of an op
rather than needing a variant for each. For an input block
dtypeis thedequantized float dtype; for
outit is the quantized storage dtype.Granularity is carried by the shape of the
scaletensor -- per-tensor,per-channel, per-group and blockwise are the same op with differently shaped
qparams.
QuantParamsStruct.is_per_channel(tensor_shape)/is_per_group(tensor_shape)take the quantized tensor's shape so a group scalesuch as
[1, K // g]over[1, K]is not mistaken for per-channel.The README writes down what is stable for backends: the argument layout and
field meanings are frozen, a new capability is a new overload or a new trailing
argument with a behavior-preserving default, and the ops are never serialized.
Fake kernels report the dtype the eager kernel actually returns, including for a
partially quantized op (quantized input, float output), where they previously
reported the int8 storage dtype.
Quantizer (
quantizer/)A torchao PT2E
Quantizerthat both annotates and fuses:OpQuantizerannotates one ATen op type and folds matched instances into afused_quantop, given aQuantizationConfigand the schema names of itsquantized activations, weights and passthrough operands.
NoopQuantizerclaims nodes with an empty annotation so a later quantizerskips them, leaving the op in float.
MaxPoolQuantizershares the input observer with output 0 ofmax_pool2d_with_indices.FusedQuantQuantizercomposes them. Annotation is first-match-wins, so anarrower quantizer is listed before the general one it overrides.
Annotations are marked
_annotated, the flag other PT2E quantizers (XNNPACK,Arm) check before claiming a node, and
fuse()only folds nodes the samequantizer annotated -- never one claimed by another quantizer.
preserved_ops()reports which ATen ops the quantizer matches, so the caller canhold exactly those back from decomposition instead of maintaining a hardcoded
list.
quantizer/qspecs.pybuilds per-tensor, per-channel and blockwise staticspecs (dynamic qspecs are not offered, since fusion needs constant qparams);
quantizer/defaults.pyprovides a default op set, withget_core_quantizersfor backends that lack per-channel support or the wider op set. Weights default
to symmetric in both
make_qconfigandmake_fused_quant_quantizer. Atransposed convolution gets its per-channel weight scales on axis 1, where its
output channels are, or per-tensor weights when
groups > 1.Coloring (
colorer.py)ColorerBaseis the step between optimization and lowering: a backend declareswhich fused ops it can run (
supported_ops) and whether a given instance iswithin its limits (
is_legal), and the base class tags the rest withdelegation_tag. Colorers compose by priority: the first to claim a node keepsit, and a later colorer skips already-tagged nodes.
The tag is a claim of ownership and the generic passes honour it -- neither
ToChannelsFirstnorDecomposeFusedQuantwill touch a colored node. So thepipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A colored node the partitioner then rejects is caught by running
DecomposeFusedQuantagain afterto_backend.A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels.
test_compile.pyhas that written out end to end.Passes
pre_quantize_passes/runs on the traced program before annotation: batch-normfolding, replacing leftover batch norms with convolutions, and recomposing
add(mask) + softmaxinto_masked_softmaxso the very-negative mask sentinelnever enters the quantized domain.
optimization_passes/runs on the fused edge graph: layout canonicalization,activation hoisting, folding add/mul into linear, quantization absorption,
embedding prequantization, linear splitting, requantize formation.
IterativePassGrouplogs a warning when a group hits its step cap.decompose_fused_quant.pyis the inverse: it turns any fused op back intodequantize + ATen + quantize. It is schema-driven, so it covers every
registered fused op rather than a fixed list, and handles multi-output ops
through
getitem. It validates every node before rewriting any, so anunsupported one (grouped/blockwise qparams) fails the pass without leaving a
partly decomposed graph. Per-channel qparams are emitted as 1-D constants
rather than a
view_copyof the full-rank one, and the qparam constants leftdead are removed.
Shared passes moved to
backends/transformsThe optimization group no longer depends on any backend. It used 22 passes from
backends/cadence/aot; none were Cadence-specific in what they do:FuseCascadedTransposeOrPermuteOps,FuseCascadedViewOps,ReplaceNopTransposeOrPermuteWithViewPass) were already empty subclasses of abackends/transformspass, which is now imported directly.backends/transforms, one module and target each. Theones that listed Cadence's q/dq ops alongside
quantized_decomposed's nowtake those op sets as class attributes holding only the standard ops.
Cadence keeps each pass under its old name as a thin subclass that adds its
own ops back, so Cadence's behavior is unchanged.
FuseQuantDequantToRequantizePassneeds a requantize op to fuse a pair whoseqparams differ, and there is no backend-neutral one. The shared pass takes a
requantize_op(defaultNone) and defaults toallow_requantize=False;Cadence sets its op and keeps
True.permute_pass_utilsgainsget_overload_packetandget_placeholders, andquantize_reorder_utilsholds the two op sets the reorder passes share.Tests no longer import Cadence either:
test/helpers.tracereplacescadence.aot.compiler.trace, and the unusedops_registrationsimports are gone.buck2 uqueryfinds nobackends/cadencetarget in the dependency closure ofany
fused_quanttarget.Pipeline (
frontend.py,passes.py,compile.py)compile_to_fused_quantruns trace -> prepare -> calibrate -> convert -> fuse ->export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
DecomposeFusedQuant.trace(..., is_qat=True)exports in training mode and keeps batch norms, forprepare_qat_pt2e; by default it exports in eval mode with batch norms folded.Either way the caller's model is left in the mode it was passed in.
The pass list is a parameter, defaulting to
get_fused_quant_passes(). Callersthat want a different pipeline pass their own, and an empty list skips
optimization entirely.
One wrinkle worth knowing when writing a backend:
DecomposeFusedQuantre-emitsthe ATen ops that were held back from decomposition, and several of them
(
aten.linear,aten.silu, ...) are not core ATen, so the edge verifier rejectsthem.
compile.pyfeedsquantizer.preserved_ops()intoEdgeCompileConfig._core_aten_ops_exception_list, which is exact byconstruction: the ops held back are precisely the ops re-emitted.
Known gaps
torchao, which is not yet a core ExecuTorch dependency.DecomposeFusedQuantrejects grouped/blockwise qparams rather than emittingdequantize_affine, so an unclaimed blockwise op fails the compile.nor documented yet.
passes.pylooks upexir_ops.edge.fused_quant.*at module scope.pipeline into opt-in groups (which takes
ToChannelsLastout of the default),a colored-node
Partitioneradapter, fitting the quantizer and frontend intoComposableQuantizer/export()recipes, a QAT fake-quant path, generatingschemas from the ATen ones, and the core/transforms/quantizer placement split.
Differential Revision: D119511656