Make QuantAbsorption's input-side absorption opt-in (#24780)

Summary:

This should have been part of
#24720, which added
`QuantAbsorptionPass` with the other fused_quant graph passes. It was left out
of that PR by mistake.

By default, `QuantAbsorptionPass` now only runs the rewrite that is never less
accurate. The docstring explains why the other two modes are opt-in.

- Output side (unchanged, always on): the producer writes its output directly
  in the downstream quantize's qparams. The value still ends up on that grid,
  rounded once instead of twice, so each element is at least as close to the
  float result. It isn't bit-exact with the q/dq graph, though.
- Input side (now opt-in, `absorb_into_inputs`): the consumer reads values on
  the dequantize's grid instead of the quantize's. Dropping the second rounding
  is more accurate on average, but not for every value.
- Forked dequantize (already opt-in, `absorb_with_fork`): its float users see
  values on the quantize's grid instead of the producer's, which can be
  coarser. Only the docstring changes.

With input-side absorption off, a pair that only the input side would remove
is left to `ReplaceDequantQuantWithRequantize`, which turns it into a
`fused_quant.requantize`.

Differential Revision: D130466673

Author: mcremon-meta

backends/transforms/fused_quant/optimization_passes/quant_absorption.py

@@ -567,11 +567,16 @@ class QuantAbsorptionPass(ExportedProgramPassBase):
     feeds into a dequant→quant pair, this pass composes the producer, dequantize,
     and downstream quantize qparams into the fused op's output qparams. When the
     producer and dequantize qparams match, this reduces to the downstream target.
+    This is always on because it is never less accurate: the value still ends up
+    on the downstream quantize's grid, rounded once instead of twice, so each
+    element is at least as close to the float result. It is not bit-exact with
+    the q/dq graph, though, and a QAT model was trained with both roundings.
 
     On the input side, when a dequant→quant pair feeds any fused_quant input with
     per-tensor qparams, this pass composes the dequantize, quantize, and consumer
     input qparams, then bypasses the pair. This handles graph inputs, which have
-    no fused producer into which the quantization can be absorbed.
+    no fused producer into which the quantization can be absorbed. It is opt-in
+    (``absorb_into_inputs``) because it is not always more accurate: see below.
 
     It also absorbs a *bare* downstream quantize into a fused op whose output is
     currently unquantized (null out qparams), turning that op into a directly
@@ -581,10 +586,14 @@ class QuantAbsorptionPass(ExportedProgramPassBase):
     Args:
         absorb_with_fork: When True, also absorb dequant→quant pairs where
             the dequant has additional users (a "forked" dequant). The dequant
-            is kept for the other users but its qparams are updated to match
-            the new output qparams. This changes (not necessarily worsens) the
-            numerics of the float path through the dequant, but should improve
-            numerics on the int path since it removes a redundant dq-q pair.
+            is kept for the other users, with the downstream quantize's qparams.
+            Off by default: those float users then see values on the quantize's
+            grid instead of the producer's, which can be coarser.
+        absorb_into_inputs: When True, also absorb a dequant→quant pair into
+            the input qparams of the fused op it feeds. Off by default: the
+            consumer then reads values on the dequantize's grid instead of the
+            quantize's. Dropping the second rounding is more accurate on
+            average, but not for every value.
         symmetric_quant_input_ops: Ops whose inputs must stay symmetrically
             quantized (zero point 0) because their backend kernel ignores input
             zero points. An input absorption that would give them a nonzero input
@@ -594,10 +603,12 @@ class QuantAbsorptionPass(ExportedProgramPassBase):
     def __init__(
         self,
         absorb_with_fork: bool = False,
+        absorb_into_inputs: bool = False,
         symmetric_quant_input_ops: frozenset[EdgeOpOverload] = frozenset(),
     ) -> None:
         super().__init__()
         self._absorb_with_fork = absorb_with_fork
+        self._absorb_into_inputs = absorb_into_inputs
         self._symmetric_quant_input_ops = symmetric_quant_input_ops
 
     def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
@@ -615,6 +626,8 @@ def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
                 modified |= _try_absorb_quant_into_unquantized_out(
                     exported_program, node, graph
                 )
+            if not self._absorb_into_inputs:
+                continue
             assert isinstance(node.target, EdgeOpOverload)
             input_names, _ = split_fused_arg_names(node.target)
             keep_symmetric = node.target in self._symmetric_quant_input_ops

backends/transforms/fused_quant/test/test_quant_absorption.py

@@ -1416,10 +1416,17 @@ def _build_dequant_quant_add(self) -> tuple[ExportedProgram, torch.Tensor]:
         builder.output([add])
         return builder.get_program(), inp
 
+    def test_does_not_absorb_into_consumer_input_by_default(self) -> None:
+        ep, _ = self._build_dequant_quant_pool()
+
+        result = QuantAbsorptionPass()(ep)
+
+        self.assertFalse(result.modified)
+
     def test_absorbs_dequant_quant_into_fused_consumer_input(self) -> None:
         ep, inp = self._build_dequant_quant_pool()
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertTrue(result.modified)
         (actual,) = result.exported_program.module()(inp)
@@ -1458,7 +1465,7 @@ def test_absorbs_dequant_quant_into_fused_consumer_input(self) -> None:
     def test_absorbs_consumer_input_through_view_copy(self) -> None:
         ep, inp = self._build_dequant_quant_pool(input_passthrough="view")
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertTrue(result.modified)
         (actual,) = result.exported_program.module()(inp)
@@ -1501,7 +1508,7 @@ def test_absorbs_consumer_input_through_view_copy(self) -> None:
     def test_absorbs_consumer_input_through_permute_copy(self) -> None:
         ep, inp = self._build_dequant_quant_pool(input_passthrough="permute")
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertTrue(result.modified)
         (actual,) = result.exported_program.module()(inp)
@@ -1546,7 +1553,7 @@ def test_does_not_absorb_through_shared_input_passthrough(self) -> None:
             input_passthrough="view", expose_passthrough=True
         )
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertFalse(result.modified)
         self.assertEqual(
@@ -1562,7 +1569,7 @@ def test_does_not_absorb_through_shared_input_passthrough(self) -> None:
     def test_absorbs_all_quantized_inputs_of_binary_op(self) -> None:
         ep, inp = self._build_dequant_quant_add()
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertTrue(result.modified)
         (actual,) = result.exported_program.module()(inp)
@@ -1628,7 +1635,7 @@ def test_composes_mismatched_consumer_input_qparams(self) -> None:
         )
         expected = ep.module()(inp)
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertTrue(result.modified)
         actual = result.exported_program.module()(inp)
@@ -1668,7 +1675,7 @@ def test_does_not_absorb_unrepresentable_consumer_input_qparams(self) -> None:
             consumer_input_zero_point=-4,
         )
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertFalse(result.modified)
 
@@ -1683,9 +1690,10 @@ def test_symmetric_quant_input_op_keeps_requantize_for_nonzero_zero_point(
         )
 
         result = QuantAbsorptionPass(
+            absorb_into_inputs=True,
             symmetric_quant_input_ops=frozenset(
                 {exir_ops.edge.fused_quant.avg_pool2d_channels_last.default}
-            )
+            ),
         )(ep)
 
         self.assertFalse(result.modified)
@@ -1696,9 +1704,10 @@ def test_symmetric_quant_input_op_absorbs_when_zero_point_stays_zero(
         ep, inp = self._build_dequant_quant_pool()
 
         result = QuantAbsorptionPass(
+            absorb_into_inputs=True,
             symmetric_quant_input_ops=frozenset(
                 {exir_ops.edge.fused_quant.avg_pool2d_channels_last.default}
-            )
+            ),
         )(ep)
 
         self.assertTrue(result.modified)
@@ -1746,7 +1755,7 @@ def test_does_not_absorb_per_channel_consumer_input(self) -> None:
         builder.output([relu])
         ep = builder.get_program()
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertFalse(result.modified)
         self.assertEqual(
@@ -1762,7 +1771,7 @@ def test_does_not_absorb_per_channel_consumer_input(self) -> None:
     def test_does_not_absorb_shared_quantize(self) -> None:
         ep, _ = self._build_dequant_quant_pool(expose_quant=True)
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertFalse(result.modified)
         self.assertEqual(
@@ -1779,7 +1788,7 @@ def test_absorbs_quantize_and_preserves_shared_dequantize(self) -> None:
         ep, inp = self._build_dequant_quant_pool(expose_dequant=True)
         _, expected_dequant = ep.module()(inp)
 
-        result = QuantAbsorptionPass()(ep)
+        result = QuantAbsorptionPass(absorb_into_inputs=True)(ep)
 
         self.assertTrue(result.modified)
         actual_pool, actual_dequant = result.exported_program.module()(inp)