Add support to sink input dequants through transparent ops (#18504)

Differential Revision: D97877439

Pull Request resolved: #18504

Author: ethansfng

backends/cadence/aot/compiler.py

@@ -239,6 +239,11 @@ def quantize_pt2(
 
     # Apply quant fusion to the exported program
     program = torch.export.export(converted_gm, inputs, strict=True)
+
+    # Sink dequant nodes through transparent ops so they fuse per-branch.
+    if quant_input_args is not None:
+        QuantizedInputWrapper.sink_dequants(program)
+
     fused_program = apply_pre_edge_transform_passes(program, quantizer)
 
     if dump_graphs:

backends/cadence/aot/compiler_funcs.py

@@ -7,6 +7,7 @@
 # pyre-strict
 
 import logging
+import operator
 from typing import Any, Optional, Union
 
 import torch
@@ -303,7 +304,26 @@ def __init__(
         if isinstance(input_args, list):
             self.quant_args = extract_input_quant_params_from_graph(module, input_args)
         elif isinstance(input_args, dict):
-            self.quant_args = input_args
+            # dict[int, QuantArgs] — use directly
+            # dict[int, list[str]] — extract quant params from graph, keyed by input index
+            first_value = next(iter(input_args.values()), None)
+            if (
+                isinstance(first_value, (list, tuple))
+                and first_value
+                and isinstance(first_value[0], str)
+            ):
+                # Values are lists of node names: extract quant params and map
+                # to the caller-specified input indices.
+                for input_idx, node_names in input_args.items():
+                    assert isinstance(node_names, list)
+                    extracted = extract_input_quant_params_from_graph(
+                        module, node_names
+                    )
+                    # Use the first extracted quant params for this input index.
+                    if extracted:
+                        self.quant_args[int(input_idx)] = next(iter(extracted.values()))
+            else:
+                self.quant_args = {int(k): v for k, v in input_args.items()}
 
     def forward(self, *args: torch.Tensor) -> Any:
         """Run inference, dequantizing configured inputs."""
@@ -349,6 +369,27 @@ def forward(self, *args: torch.Tensor) -> Any:
 
         return self.module(*dequantized_args)
 
+    @staticmethod
+    def sink_dequants(program: torch.export.ExportedProgram) -> None:
+        """Sink dequant nodes through transparent ops in an exported program.
+
+        If the graph branches through transparent ops (view, split, getitem, etc.)
+        into paths with different quantization parameters, sink the dequants to be
+        adjacent to each downstream quant node, enabling per-branch fusion.
+
+        Must be called after export() on a QuantizedInputWrapper-wrapped model.
+        """
+        from torch.export.graph_signature import InputKind
+
+        user_input_names = {
+            spec.arg.name
+            for spec in program.graph_signature.input_specs
+            if spec.kind == InputKind.USER_INPUT
+        }
+        sink_input_dequant_through_transparent_ops(
+            program.graph_module, user_input_names
+        )
+
 
 class QuantizedOutputWrapper(torch.nn.Module):
     """
@@ -379,3 +420,183 @@ def forward(self, *args: torch.Tensor) -> Any:
         return torch.ops.quantized_decomposed.quantize_per_tensor.default(
             result, scale, zp, qmin, qmax, dtype
         )
+
+
+def _get_transparent_ops() -> set[Any]:
+    """Ops that only reshape/index data without changing values.
+    Safe to pass uint8 data through these."""
+    return {
+        torch.ops.aten.view_copy.default,
+        torch.ops.aten.view.default,
+        torch.ops.aten.reshape.default,
+        torch.ops.aten.split.Tensor,
+        torch.ops.aten.slice_copy.Tensor,
+        torch.ops.aten.permute_copy.default,
+        torch.ops.aten.permute.default,
+        torch.ops.aten.expand_copy.default,
+        torch.ops.aten.unsqueeze_copy.default,
+        torch.ops.aten.squeeze_copy.dim,
+        torch.ops.aten.transpose_copy.int,
+        torch.ops.aten.clone.default,
+        operator.getitem,
+    }
+
+
+def _get_quantize_ops() -> set[Any]:
+    ops = {torch.ops.quantized_decomposed.quantize_per_tensor.default}
+    try:
+        ops.add(torch.ops.cadence.quantize_per_tensor.default)
+    except AttributeError:
+        pass
+    return ops
+
+
+def _get_dequantize_ops() -> set[Any]:
+    ops = {torch.ops.quantized_decomposed.dequantize_per_tensor.default}
+    try:
+        ops.add(torch.ops.cadence.dequantize_per_tensor.default)
+    except AttributeError:
+        pass
+    return ops
+
+
+def _walk_to_downstream_quants(
+    node: torch.fx.Node,
+    quantize_ops: set[Any],
+    transparent_ops: set[Any],
+    downstream_quants: list[torch.fx.Node],
+) -> bool:
+    """Walk forward through transparent ops collecting downstream quant nodes.
+
+    Returns True if all paths end at a quant node.
+    """
+    all_valid = True
+    for user in node.users:
+        if user.op == "call_function" and user.target in quantize_ops:
+            downstream_quants.append(user)
+        elif user.op == "call_function" and user.target in transparent_ops:
+            if not _walk_to_downstream_quants(
+                user, quantize_ops, transparent_ops, downstream_quants
+            ):
+                all_valid = False
+        else:
+            all_valid = False
+    return all_valid
+
+
+def _get_dequant_node_for_placeholder(
+    placeholder: torch.fx.Node,
+    input_placeholder_names: set[str] | None,
+    dequantize_ops: set[Any],
+) -> torch.fx.Node | None:
+    """Return the single dequant user of a uint8 placeholder, or None."""
+    if placeholder.op != "placeholder":
+        return None
+    if (
+        input_placeholder_names is not None
+        and placeholder.name not in input_placeholder_names
+    ):
+        return None
+    val = placeholder.meta.get("val")
+    if val is None or not isinstance(val, torch.Tensor):
+        return None
+    if val.dtype != torch.uint8:
+        return None
+    if len(placeholder.users) != 1:
+        return None
+    dequant_node = next(iter(placeholder.users))
+    if dequant_node.op == "call_function" and dequant_node.target in dequantize_ops:
+        return dequant_node
+    return None
+
+
+def _sink_dequant_to_quant_nodes(
+    graph: torch.fx.Graph,
+    dequant_node: torch.fx.Node,
+    placeholder: torch.fx.Node,
+    downstream_quants: list[torch.fx.Node],
+) -> None:
+    """Insert per-branch dequants before each downstream quant and rewire."""
+    dequant_op = dequant_node.target
+    assert callable(dequant_op)
+
+    for quant_node in downstream_quants:
+        quant_input = quant_node.args[0]
+        assert isinstance(quant_input, torch.fx.Node)
+        quant_params = quant_node.args[1:]
+
+        with graph.inserting_before(quant_node):
+            new_dequant = graph.call_function(
+                dequant_op,
+                args=(quant_input, *quant_params),
+            )
+        new_dequant.meta = {**dequant_node.meta}
+        if "val" in quant_node.meta and isinstance(
+            quant_node.meta["val"], torch.Tensor
+        ):
+            quant_val = quant_node.meta["val"]
+            new_dequant.meta["val"] = torch.empty(quant_val.shape, dtype=torch.float32)
+
+        quant_node.replace_input_with(quant_input, new_dequant)
+
+    dequant_node.replace_all_uses_with(placeholder)
+    graph.erase_node(dequant_node)
+
+
+def sink_input_dequant_through_transparent_ops(
+    graph_module: GraphModule,
+    input_placeholder_names: set[str] | None = None,
+) -> bool:
+    """
+    Sinks dequantize nodes from quantized input placeholders through transparent ops
+    to be adjacent to downstream quantize nodes, enabling dequant-quant fusion.
+    This creates per-branch dequants with matching params.
+
+    Args:
+        graph_module: The graph module to transform.
+        input_placeholder_names: Optional set of placeholder names to consider.
+            If provided, only these placeholders are processed (use this to
+            restrict to user inputs and avoid touching weight/buffer placeholders).
+            If None, all uint8 placeholders are considered.
+
+    Returns True if the graph was modified.
+    """
+    graph = graph_module.graph
+    modified = False
+
+    transparent_ops: set[Any] = _get_transparent_ops()
+    quantize_ops: set[Any] = _get_quantize_ops()
+    dequantize_ops: set[Any] = _get_dequantize_ops()
+
+    for placeholder in list(graph.nodes):
+        dequant_node = _get_dequant_node_for_placeholder(
+            placeholder, input_placeholder_names, dequantize_ops
+        )
+        if dequant_node is None:
+            continue
+
+        downstream_quants: list[torch.fx.Node] = []
+        all_paths_end_at_quant = _walk_to_downstream_quants(
+            dequant_node, quantize_ops, transparent_ops, downstream_quants
+        )
+
+        if not downstream_quants or not all_paths_end_at_quant:
+            continue
+
+        _sink_dequant_to_quant_nodes(
+            graph, dequant_node, placeholder, downstream_quants
+        )
+
+        modified = True
+        logger.info(
+            "Sunk dequant for input '%s' through transparent ops to %d "
+            "downstream quant nodes",
+            placeholder.name,
+            len(downstream_quants),
+        )
+
+    if modified:
+        graph.lint()
+        graph_module.recompile()
+
+    return modified