Add QuantizedOutputWrapper API

Differential Revision: D92579429

Pull Request resolved: #17289

Author: ethansfng

backends/cadence/aot/compiler_funcs.py

@@ -19,29 +19,6 @@
 QuantArgs = tuple[float, int, int, int, torch.dtype]
 
 
-def extract_input_shapes_from_graph(
-    module: GraphModule,
-) -> dict[int, tuple[int, ...]]:
-    """
-    Extract input shapes from the FX graph placeholder nodes.
-
-    Returns a dict mapping input index to expected shape tuple.
-    """
-    input_shapes: dict[int, tuple[int, ...]] = {}
-    idx = 0
-    for node in module.graph.nodes:
-        if node.op == "placeholder":
-            # Get the tensor_meta from the node if available
-            if "val" in node.meta:
-                val = node.meta["val"]
-                if isinstance(val, torch.Tensor):
-                    input_shapes[idx] = tuple(val.shape)
-                elif hasattr(val, "shape"):
-                    input_shapes[idx] = tuple(val.shape)
-            idx += 1
-    return input_shapes
-
-
 @torch.no_grad()
 def trace(
     model: torch.nn.Module,
@@ -81,6 +58,29 @@ def prepare(
     return prepared_model
 
 
+def extract_input_shapes_from_graph(
+    module: GraphModule,
+) -> dict[int, tuple[int, ...]]:
+    """
+    Extract input shapes from the FX graph placeholder nodes.
+
+    Returns a dict mapping input index to expected shape tuple.
+    """
+    input_shapes: dict[int, tuple[int, ...]] = {}
+    idx = 0
+    for node in module.graph.nodes:
+        if node.op == "placeholder":
+            # Get the tensor_meta from the node if available
+            if "val" in node.meta:
+                val = node.meta["val"]
+                if isinstance(val, torch.Tensor):
+                    input_shapes[idx] = tuple(val.shape)
+                elif hasattr(val, "shape"):
+                    input_shapes[idx] = tuple(val.shape)
+            idx += 1
+    return input_shapes
+
+
 def extract_quant_params_through_permute(
     module: torch.fx.GraphModule,
 ) -> dict[int, tuple[float, int, int, int, torch.dtype]]:
@@ -121,6 +121,83 @@ def extract_quant_params_through_permute(
     return quant_args
 
 
+def extract_output_dequant_params(
+    module: torch.fx.GraphModule,
+) -> QuantArgs:
+    """
+    Extract dequantization parameters from the output of a quantized model.
+
+    The graph is expected to end with:
+        ... → dequantize_per_tensor(scale, zp, qmin, qmax, dtype) → output
+    """
+    for node in module.graph.nodes:
+        if node.op != "output":
+            continue
+        output_args = node.args[0]
+        if isinstance(output_args, (tuple, list)):
+            target_output = output_args[0]
+        else:
+            target_output = output_args
+        if not isinstance(target_output, torch.fx.Node):
+            raise ValueError("Output node is not an FX node")
+        if "dequantize_per_tensor" in str(target_output.target):
+            args = target_output.args[1:]
+            if len(args) >= 5:
+                dtype = args[4]
+                assert isinstance(dtype, torch.dtype)
+                return (
+                    float(args[0]),  # scale
+                    int(args[1]),  # zero_point
+                    int(args[2]),  # qmin
+                    int(args[3]),  # qmax
+                    dtype,
+                )
+    raise ValueError("Could not find dequantize_per_tensor at the output of the graph")
+
+
+def extract_output_dequant_params_through_permute(
+    module: torch.fx.GraphModule,
+) -> QuantArgs:
+    """
+    Extract dequantization parameters from the output through a permute.
+
+    For models with nhwc output, the graph ends with:
+        ... → dequantize_per_tensor → permute(0, 2, 3, 1) → output
+    """
+    for node in module.graph.nodes:
+        if node.op != "output":
+            continue
+        output_args = node.args[0]
+        if isinstance(output_args, (tuple, list)):
+            target_output = output_args[0]
+        else:
+            target_output = output_args
+        if not isinstance(target_output, torch.fx.Node):
+            raise ValueError("Output node is not an FX node")
+        if target_output.target in (
+            torch.ops.aten.permute.default,
+            torch.ops.aten.permute_copy.default,
+        ):
+            permute_input = target_output.args[0]
+            if isinstance(
+                permute_input, torch.fx.Node
+            ) and "dequantize_per_tensor" in str(permute_input.target):
+                args = permute_input.args[1:]
+                if len(args) >= 5:
+                    dtype = args[4]
+                    assert isinstance(dtype, torch.dtype)
+                    return (
+                        float(args[0]),  # scale
+                        int(args[1]),  # zero_point
+                        int(args[2]),  # qmin
+                        int(args[3]),  # qmax
+                        dtype,
+                    )
+    raise ValueError(
+        "Could not find dequantize_per_tensor → permute at the output of the graph"
+    )
+
+
 def extract_input_quant_params_from_graph(
     module: GraphModule,
     input_names: list[str],
@@ -241,3 +318,34 @@ def forward(self, *args: torch.Tensor) -> Any:
             dequantized_args.append(node)
 
         return self.module(*dequantized_args)
+
+
+class QuantizedOutputWrapper(torch.nn.Module):
+    """
+    Wrapper that quantizes a model's output so it produces uint8 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) →
+    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.
+    """
+
+    def __init__(
+        self,
+        module: torch.nn.Module,
+        output_quant_args: QuantArgs,
+    ) -> None:
+        super().__init__()
+        self.module: torch.nn.Module = module
+        self.output_quant_args: QuantArgs = 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
+        )