Generalize QuantizedOutputWrapper for multi-output models

Summary: Add support for multiple outputs in quantized output wrapper

Differential Revision: D107429509

Author: ethansfng

backends/cadence/aot/compiler_funcs.py

@@ -159,6 +159,47 @@ def extract_output_dequant_params(
     raise ValueError("Could not find dequantize_per_tensor at the output of the graph")
 
 
+def extract_all_output_dequant_params(
+    module: torch.fx.GraphModule,
+) -> list[QuantArgs | None]:
+    """
+    Extract per-output dequantization parameters from a multi-output model.
+
+    Returns a QuantArgs tuple for outputs ending in dequantize_per_tensor
+    or None for outputs that aren't dequantized.
+    """
+    for node in module.graph.nodes:
+        if node.op != "output":
+            continue
+        output_args = node.args[0]
+        if not isinstance(output_args, (tuple, list)):
+            output_args = (output_args,)
+        params: list[QuantArgs | None] = []
+        for out in output_args:
+            if isinstance(out, torch.fx.Node) and "dequantize_per_tensor" in str(
+                out.target
+            ):
+                args = out.args[1:]
+                if len(args) >= 5:
+                    dtype = args[4]
+                    assert isinstance(dtype, torch.dtype)
+                    params.append(
+                        (
+                            float(args[0]),
+                            int(args[1]),
+                            int(args[2]),
+                            int(args[3]),
+                            dtype,
+                        )
+                    )
+                else:
+                    params.append(None)
+            else:
+                params.append(None)
+        return params
+    raise ValueError("No output node in graph")
+
+
 def extract_output_dequant_params_through_permute(
     module: torch.fx.GraphModule,
 ) -> QuantArgs:
@@ -400,33 +441,55 @@ def sink_dequants(program: torch.export.ExportedProgram) -> None:
 
 class QuantizedOutputWrapper(torch.nn.Module):
     """
-    Wrapper that quantizes a model's output so it produces uint8 tensors.
+    Wrapper that quantizes a model's output(s) so they produce quantized tensors.
 
     Mirrors QuantizedInputWrapper: the wrapper adds a quantize_per_tensor after
-    the model's output. When the graph is traced, the dequant (from the model) →
+    each output. When the graph is traced, the dequant (from the model) →
     quant (from the wrapper) pair with matching parameters folds away, leaving
     the output in its quantized form.
 
     Args:
         module: The module to wrap (may already be a QuantizedInputWrapper).
-        output_quant_args: (scale, zero_point, qmin, qmax, dtype) for the output.
+        output_quant_args: Quantization parameters — either a single QuantArgs
+            tuple or a list with one entry per output.
     """
 
     def __init__(
         self,
         module: torch.nn.Module,
-        output_quant_args: QuantArgs,
+        output_quant_args: Union[QuantArgs, list[QuantArgs | None]],
     ) -> None:
         super().__init__()
         self.module: torch.nn.Module = module
-        self.output_quant_args: QuantArgs = output_quant_args
+        if isinstance(output_quant_args, list):
+            self._multi_output: bool = True
+            self._per_output_args: list[QuantArgs | None] = output_quant_args
+        else:
+            self._multi_output = False
+            self._per_output_args = [output_quant_args]
 
     def forward(self, *args: torch.Tensor) -> Any:
         result = self.module(*args)
-        scale, zp, qmin, qmax, dtype = self.output_quant_args
-        return torch.ops.quantized_decomposed.quantize_per_tensor.default(
-            result, scale, zp, qmin, qmax, dtype
-        )
+        if not self._multi_output:
+            qa = self._per_output_args[0]
+            assert qa is not None
+            scale, zp, qmin, qmax, dtype = qa
+            return torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                result, scale, zp, qmin, qmax, dtype
+            )
+        out: list[torch.Tensor] = []
+        for i, r in enumerate(result):
+            qa = self._per_output_args[i] if i < len(self._per_output_args) else None
+            if qa is None:
+                out.append(r)
+                continue
+            scale, zp, qmin, qmax, dtype = qa
+            out.append(
+                torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                    r, scale, zp, qmin, qmax, dtype
+                )
+            )
+        return tuple(out)
 
 
 def _get_transparent_ops() -> set[Any]: