Add fuse() to remaining QuantizationPatterns (pytorch#19727)

Summary:

Add `fuse()` implementations to the remaining Cadence `QuantizationPattern` subclasses:

- `MaxPool2dPattern`, `MaxPool2dWithoutIndicesPattern` — order-preserving pool on quantized values
- `ReluBasePattern` (inherited by `ReluPattern0`/`1`) — relu with requantization
- `ConvReluBasePattern` (inherited by `Conv1d`/`2dReluPattern0`/`1`) — conv+relu fusion with `anchor_ops()` override to match only the conv op
- `SoftmaxPattern` — softmax with dummy mask/pos tensors and fake_mode metadata
- `MixedW8A32LinearPattern` — weight-only quantized linear (no input/output quant)
- `MixedW8A32ConvPattern` — weight-only quantized conv1d with NCL→NLC permutation
- `MixedW8A32GruPattern` — weight-only quantized GRU with 4 dequantized params

Differential Revision: D105728177

Author: ethansfng

backends/cadence/aot/quantizer/patterns.py

@@ -12,6 +12,7 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
+from executorch.backends.cadence.aot.compiler_utils import get_shape
 from executorch.backends.cadence.aot.pass_utils import get_arg
 from executorch.backends.cadence.aot.quantizer.pattern_utils import (
     DQ_PER_TENSOR,
@@ -25,6 +26,7 @@
     check_out_zero_point_is_min_range,
     copy_node_metadata,
     get_bias_qparams,
+    quantize_tensor_multiplier,
 )
 from torch import fx
 from torch._ops import OpOverload
@@ -321,9 +323,8 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
         quant_node = find_quant_user(relu_node)
         if quant_node is None:
             return None
-        check_out_zero_point_is_min_range(
-            quant_node.args[2], quant_node.args[5]  # pyre-ignore[6]
-        )
+        # pyre-ignore[6]: Argument -> int/dtype narrowing
+        check_out_zero_point_is_min_range(quant_node.args[2], quant_node.args[5])
         args = (
             dq0.args[0],
             dq0.args[1],
@@ -724,13 +725,11 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
         quant_node = find_quant_user(anchor_node)
         if quant_node is None:
             return None
-        dq_bias = (
-            anchor_node.args[2]
-            if len(anchor_node.args) > 2
-            and isinstance(anchor_node.args[2], fx.Node)
-            and anchor_node.args[2].target == DQ_PER_TENSOR
-            else None
-        )
+        dq_bias: Optional[fx.Node] = None
+        if len(anchor_node.args) > 2:
+            bias_arg = anchor_node.args[2]
+            if isinstance(bias_arg, fx.Node) and bias_arg.target == DQ_PER_TENSOR:
+                dq_bias = bias_arg
         return fuse_linear(
             gm,
             dq_input,
@@ -836,6 +835,30 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_max_pool2d_nchw.default
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        dq_input = anchor_node.args[0]
+        if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
+            return None
+        quant_node = find_quant_user(anchor_node)
+        if quant_node is None:
+            return None
+        kernel_size = get_arg(anchor_node, "kernel_size", Optional[list[int]]) or [1, 1]
+        stride = get_arg(anchor_node, "stride", Optional[list[int]]) or kernel_size
+        padding = get_arg(anchor_node, "padding", Optional[list[int]]) or [0, 0]
+        dilation = get_arg(anchor_node, "dilation", Optional[list[int]]) or [1, 1]
+        ceil_mode = get_arg(anchor_node, "ceil_mode", Optional[bool]) or False
+        args = (dq_input.args[0],)
+        kwargs = {
+            "kernel_size": kernel_size,
+            "stride": stride,
+            "padding": padding,
+            "dilation": dilation,
+            "ceil_mode": ceil_mode,
+        }
+        return replace_with_op(
+            gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
+        )
+
 
 class MaxPool2dWithoutIndicesPattern(QuantizationPattern):
     """
@@ -875,6 +898,30 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_max_pool2d_nchw.default
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        dq_input = anchor_node.args[0]
+        if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
+            return None
+        quant_node = find_quant_user(anchor_node)
+        if quant_node is None:
+            return None
+        kernel_size = get_arg(anchor_node, "kernel_size", Optional[list[int]]) or [1, 1]
+        stride = get_arg(anchor_node, "stride", Optional[list[int]]) or kernel_size
+        padding = get_arg(anchor_node, "padding", Optional[list[int]]) or [0, 0]
+        dilation = get_arg(anchor_node, "dilation", Optional[list[int]]) or [1, 1]
+        ceil_mode = get_arg(anchor_node, "ceil_mode", Optional[bool]) or False
+        args = (dq_input.args[0],)
+        kwargs = {
+            "kernel_size": kernel_size,
+            "stride": stride,
+            "padding": padding,
+            "dilation": dilation,
+            "ceil_mode": ceil_mode,
+        }
+        return replace_with_op(
+            gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
+        )
+
 
 # This is a base class for ReLU
 
@@ -904,6 +951,29 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_relu.per_tensor
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        dq_input = anchor_node.args[0]
+        if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
+            return None
+        quant_node = find_quant_user(anchor_node)
+        if quant_node is None:
+            return None
+        input_scale = dq_input.args[1]
+        # pyre-fixme[58]
+        requantize_scale = input_scale / quant_node.args[1]
+        requantize_scale_t = torch.tensor([requantize_scale])
+        out_multiplier, out_shift = quantize_tensor_multiplier(requantize_scale_t)
+        args = (dq_input.args[0],)
+        kwargs = {
+            "X_zero_point": dq_input.args[2],
+            "out_zero_point": quant_node.args[2],
+            "out_multiplier": out_multiplier[0].item(),
+            "out_shift": out_shift[0].item(),
+        }
+        return replace_with_op(
+            gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
+        )
+
 
 # Regular relu op
 class ReluPattern0(ReluBasePattern):
@@ -963,6 +1033,37 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_conv2d_nchw.per_tensor
 
+    def anchor_ops(self) -> tuple[OpOverload, ...]:
+        return (self.partition_types()[0],)
+
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        conv_users = list(anchor_node.users)
+        if len(conv_users) != 1:
+            return None
+        relu_node = conv_users[0]
+        if relu_node.target != self.partition_types()[1]:
+            return None
+        _arg0 = anchor_node.args[0]
+        dq_input = (
+            _arg0
+            if isinstance(_arg0, fx.Node) and _arg0.target == DQ_PER_TENSOR
+            else None
+        )
+        _arg1 = anchor_node.args[1]
+        dq_weight = (
+            _arg1
+            if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR
+            else None
+        )
+        if dq_input is None or dq_weight is None:
+            return None
+        quant_node = find_quant_user(relu_node)
+        if quant_node is None:
+            return None
+        # pyre-ignore[6]: Argument -> int/dtype narrowing
+        check_out_zero_point_is_min_range(quant_node.args[2], quant_node.args[5])
+        return fuse_conv(self, gm, anchor_node, dq_input, dq_weight, quant_node)
+
 
 # Conv1d + regular relu op fusion
 class Conv1dReluPattern0(ConvReluBasePattern):
@@ -1017,6 +1118,52 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_softmax.per_tensor
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        dq_input = anchor_node.args[0]
+        if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
+            return None
+        quant_node = find_quant_user(anchor_node)
+        if quant_node is None:
+            return None
+        input_q = dq_input.args[0]
+        assert isinstance(input_q, fx.Node)
+        quant_input = quant_node.args[0]
+        assert isinstance(quant_input, fx.Node)
+        mask_shape = get_shape(gm, quant_input)
+        mask_shape = list(mask_shape) if mask_shape else []
+        mask_shape[-1] = mask_shape[-1] // 16
+        with gm.graph.inserting_before(anchor_node):
+            mask_tensor = gm.graph.call_function(
+                torch.ops.aten.full.default, (mask_shape, 0.0), {"dtype": torch.int32}
+            )
+        assert "val" in input_q.meta
+        fake_mode = input_q.meta["val"].fake_mode
+        assert fake_mode is not None
+        with fake_mode:
+            mask_tensor.meta["val"] = torch.full(mask_shape, 0.0, dtype=torch.int32)
+        copy_node_metadata(mask_tensor, input_q)
+        with gm.graph.inserting_before(anchor_node):
+            pos_tensor = gm.graph.call_function(
+                torch.ops.aten.full.default, ([1], 0), {"dtype": torch.int64}
+            )
+        with fake_mode:
+            pos_tensor.meta["val"] = torch.full([1], 0, dtype=torch.int64)
+        copy_node_metadata(pos_tensor, input_q)
+        args = (
+            input_q,
+            mask_tensor,
+            get_arg(anchor_node, "dim", int),
+            0,
+            pos_tensor,
+            dq_input.args[1],
+            dq_input.args[2],
+            quant_node.args[1],
+            quant_node.args[2],
+        )
+        return replace_with_op(
+            gm, anchor_node, self.replacement_op(), args, {}, quant_node
+        )
+
 
 class MixedW8A32LinearPattern(QuantizationPattern):
     def partition_types(self) -> List[OpOverload]:
@@ -1071,6 +1218,38 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_w8a32_linear.default
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        if len(anchor_node.args) != 3 or len(anchor_node.kwargs) > 0:
+            return None
+        _arg1 = anchor_node.args[1]
+        dq_weight = (
+            _arg1
+            if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR
+            else None
+        )
+        _arg2 = anchor_node.args[2]
+        dq_bias = (
+            _arg2
+            if isinstance(_arg2, fx.Node) and _arg2.target == DQ_PER_TENSOR
+            else None
+        )
+        if dq_weight is None or dq_bias is None:
+            return None
+        input_node = anchor_node.args[0]
+        assert isinstance(input_node, fx.Node)
+        args = (
+            input_node,
+            dq_weight.args[0],
+            dq_weight.args[1],
+            dq_bias.args[0],
+            dq_bias.args[1],
+        )
+        with gm.graph.inserting_after(anchor_node):
+            fused = gm.graph.call_function(self.replacement_op(), args, {})
+        fused.meta = anchor_node.meta
+        anchor_node.replace_all_uses_with(fused)
+        return fused
+
 
 class MixedW8A32ConvPattern(QuantizationPattern):
     def partition_types(self) -> List[OpOverload]:
@@ -1145,6 +1324,78 @@ def get_anchors(
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_w8a32_conv.default
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        if len(anchor_node.args) != 3 or len(anchor_node.kwargs) > 0:
+            return None
+        _arg1 = anchor_node.args[1]
+        dq_weight = (
+            _arg1
+            if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR
+            else None
+        )
+        _arg2 = anchor_node.args[2]
+        dq_bias = (
+            _arg2
+            if isinstance(_arg2, fx.Node) and _arg2.target == DQ_PER_TENSOR
+            else None
+        )
+        if dq_weight is None or dq_bias is None:
+            return None
+        input_node = anchor_node.args[0]
+        assert isinstance(input_node, fx.Node)
+        assert get_arg(anchor_node, "stride", list[int]) == [1]
+        assert get_arg(anchor_node, "padding", list[int]) == [0]
+        assert get_arg(anchor_node, "dilation", list[int]) == [1]
+        assert get_arg(anchor_node, "groups", int) == 1
+        weight_q = dq_weight.args[0]
+        assert isinstance(weight_q, fx.Node)
+        with gm.graph.inserting_before(anchor_node):
+            transposed_inputs = gm.graph.call_function(
+                torch.ops.aten.permute.default, (input_node, [0, 2, 1])
+            )
+        if "val" in input_node.meta:
+            original_val = input_node.meta["val"]
+            fake_mode = original_val.fake_mode
+            if fake_mode is not None:
+                with fake_mode:
+                    transposed_inputs.meta["val"] = torch.ops.aten.permute.default(
+                        original_val, [0, 2, 1]
+                    )
+            else:
+                transposed_inputs.meta["val"] = torch.ops.aten.permute.default(
+                    original_val, [0, 2, 1]
+                )
+        copy_node_metadata(transposed_inputs, input_node)
+        with gm.graph.inserting_before(anchor_node):
+            transposed_weights = gm.graph.call_function(
+                torch.ops.aten.permute.default, (weight_q, [2, 0, 1])
+            )
+        if "val" in weight_q.meta:
+            original_val = weight_q.meta["val"]
+            fake_mode = original_val.fake_mode
+            if fake_mode is not None:
+                with fake_mode:
+                    transposed_weights.meta["val"] = torch.ops.aten.permute.default(
+                        original_val, [2, 0, 1]
+                    )
+            else:
+                transposed_weights.meta["val"] = torch.ops.aten.permute.default(
+                    original_val, [2, 0, 1]
+                )
+        copy_node_metadata(transposed_weights, weight_q)
+        args = (
+            transposed_inputs,
+            transposed_weights,
+            dq_weight.args[1],
+            dq_bias.args[0],
+            dq_bias.args[1],
+        )
+        with gm.graph.inserting_after(anchor_node):
+            fused = gm.graph.call_function(self.replacement_op(), args, {})
+        fused.meta = anchor_node.meta
+        anchor_node.replace_all_uses_with(fused)
+        return fused
+
 
 class MixedW8A32GruPattern(QuantizationPattern):
     def partition_types(self) -> List[OpOverload]:
@@ -1217,6 +1468,43 @@ def __init__(self, args, meta):
     def replacement_op(self) -> OpOverload:
         return torch.ops.cadence.quantized_w8a32_gru.default
 
+    def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
+        if len(anchor_node.kwargs) > 0:
+            return None
+        params = anchor_node.args[2]
+        if not isinstance(params, (list, tuple)) or len(params) < 4:
+            return None
+        dq_w_ih = params[0]
+        if not isinstance(dq_w_ih, fx.Node) or dq_w_ih.target != DQ_PER_TENSOR:
+            return None
+        dq_w_hh = params[1]
+        if not isinstance(dq_w_hh, fx.Node) or dq_w_hh.target != DQ_PER_TENSOR:
+            return None
+        dq_b_ih = params[2]
+        if not isinstance(dq_b_ih, fx.Node) or dq_b_ih.target != DQ_PER_TENSOR:
+            return None
+        dq_b_hh = params[3]
+        if not isinstance(dq_b_hh, fx.Node) or dq_b_hh.target != DQ_PER_TENSOR:
+            return None
+        input_node = anchor_node.args[0]
+        hidden_node = anchor_node.args[1]
+        args = (
+            input_node,
+            hidden_node,
+            dq_w_ih.args[0],
+            dq_w_ih.args[1],
+            dq_w_hh.args[0],
+            dq_w_hh.args[1],
+            dq_b_ih.args[0],
+            dq_b_ih.args[1],
+            dq_b_hh.args[0],
+        )
+        with gm.graph.inserting_after(anchor_node):
+            fused = gm.graph.call_function(self.replacement_op(), args, {})
+        fused.meta = anchor_node.meta
+        anchor_node.replace_all_uses_with(fused)
+        return fused
+
 
 class RmsNormPattern(QuantizationPattern):
     """Pattern that preserves rms_norm from decomposition without matching anything."""