Skip to content

Commit 885ee69

Browse files
ethansfngfacebook-github-bot
ethansfng
authored and
facebook-github-bot
committed
Add fuse() to remaining QuantizationPatterns (#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 Reviewed By: DrJessop Differential Revision: D105728177
1 parent 446edf8 commit 885ee69

1 file changed

Lines changed: 260 additions & 2 deletions

File tree

‎backends/cadence/aot/quantizer/patterns.py‎

Lines changed: 260 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from typing import List, Tuple, Union
1313

1414
import torch
15+
from executorch.backends.cadence.aot.compiler_utils import get_shape
1516
from executorch.backends.cadence.aot.pass_utils import get_arg, replace_with_op
1617
from executorch.backends.cadence.aot.quantizer.pattern_utils import (
1718
DQ_PER_TENSOR,
@@ -24,6 +25,7 @@
2425
from executorch.backends.cadence.aot.quantizer.utils import (
2526
check_out_zero_point_is_min_range,
2627
get_bias_qparams,
28+
quantize_tensor_multiplier,
2729
)
2830
from torch import fx
2931
from torch._ops import OpOverload
@@ -804,6 +806,40 @@ def get_anchors(
804806
def replacement_op(self) -> OpOverload:
805807
return torch.ops.cadence.quantized_max_pool2d_nchw.default
806808

809+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
810+
return _fuse_max_pool2d(gm, anchor_node)
811+
812+
813+
def _fuse_max_pool2d(gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
814+
"""Shared fuse logic for both MaxPool2d variants."""
815+
dq_input = anchor_node.args[0]
816+
if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
817+
return None
818+
quant_node = find_quant_user(anchor_node)
819+
if quant_node is None:
820+
return None
821+
kernel_size = get_arg(anchor_node, "kernel_size", list[int])
822+
stride = get_arg(anchor_node, "stride", list[int])
823+
padding = get_arg(anchor_node, "padding", list[int])
824+
dilation = get_arg(anchor_node, "dilation", list[int])
825+
ceil_mode = get_arg(anchor_node, "ceil_mode", bool)
826+
args = (get_arg(dq_input, "input", fx.Node),)
827+
kwargs = {
828+
"kernel_size": kernel_size,
829+
"stride": stride,
830+
"padding": padding,
831+
"dilation": dilation,
832+
"ceil_mode": ceil_mode,
833+
}
834+
return replace_with_op(
835+
gm,
836+
anchor_node,
837+
torch.ops.cadence.quantized_max_pool2d_nchw.default,
838+
args,
839+
kwargs,
840+
quant_node,
841+
)
842+
807843

808844
class MaxPool2dWithoutIndicesPattern(QuantizationPattern):
809845
"""
@@ -843,8 +879,8 @@ def get_anchors(
843879
def replacement_op(self) -> OpOverload:
844880
return torch.ops.cadence.quantized_max_pool2d_nchw.default
845881

846-
847-
# This is a base class for ReLU
882+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
883+
return _fuse_max_pool2d(gm, anchor_node)
848884

849885

850886
# This is a base class for ReLU, since it can be used with two different aten ops
@@ -872,6 +908,28 @@ def get_anchors(
872908
def replacement_op(self) -> OpOverload:
873909
return torch.ops.cadence.quantized_relu.per_tensor
874910

911+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
912+
dq_input = anchor_node.args[0]
913+
if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
914+
return None
915+
quant_node = find_quant_user(anchor_node)
916+
if quant_node is None:
917+
return None
918+
input_scale = get_arg(dq_input, "scale", float)
919+
requantize_scale = input_scale / get_arg(quant_node, "scale", float)
920+
requantize_scale_t = torch.tensor([requantize_scale])
921+
out_multiplier, out_shift = quantize_tensor_multiplier(requantize_scale_t)
922+
args = (get_arg(dq_input, "input", fx.Node),)
923+
kwargs = {
924+
"X_zero_point": get_arg(dq_input, "zero_point", int),
925+
"out_zero_point": get_arg(quant_node, "zero_point", int),
926+
"out_multiplier": out_multiplier[0].item(),
927+
"out_shift": out_shift[0].item(),
928+
}
929+
return replace_with_op(
930+
gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
931+
)
932+
875933

876934
# Regular relu op
877935
class ReluPattern0(ReluBasePattern):
@@ -931,6 +989,39 @@ def get_anchors(
931989
def replacement_op(self) -> OpOverload:
932990
return torch.ops.cadence.quantized_conv2d_nchw.per_tensor
933991

992+
def anchor_ops(self) -> tuple[OpOverload, ...]:
993+
return (self.partition_types()[0],)
994+
995+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
996+
conv_users = list(anchor_node.users)
997+
if len(conv_users) != 1:
998+
return None
999+
relu_node = conv_users[0]
1000+
if relu_node.target != self.partition_types()[1]:
1001+
return None
1002+
_arg0 = anchor_node.args[0]
1003+
dq_input = (
1004+
_arg0
1005+
if isinstance(_arg0, fx.Node) and _arg0.target == DQ_PER_TENSOR
1006+
else None
1007+
)
1008+
_arg1 = anchor_node.args[1]
1009+
dq_weight = (
1010+
_arg1
1011+
if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR
1012+
else None
1013+
)
1014+
if dq_input is None or dq_weight is None:
1015+
return None
1016+
quant_node = find_quant_user(relu_node)
1017+
if quant_node is None:
1018+
return None
1019+
check_out_zero_point_is_min_range(
1020+
get_arg(quant_node, "zero_point", int),
1021+
get_arg(quant_node, "dtype", torch.dtype),
1022+
)
1023+
return fuse_conv(self, gm, anchor_node, dq_input, dq_weight, quant_node)
1024+
9341025

9351026
# Conv1d + regular relu op fusion
9361027
class Conv1dReluPattern0(ConvReluBasePattern):
@@ -985,6 +1076,56 @@ def get_anchors(
9851076
def replacement_op(self) -> OpOverload:
9861077
return torch.ops.cadence.quantized_softmax.per_tensor
9871078

1079+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
1080+
dq_input = anchor_node.args[0]
1081+
if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR:
1082+
return None
1083+
quant_node = find_quant_user(anchor_node)
1084+
if quant_node is None:
1085+
return None
1086+
input_q = get_arg(dq_input, "input", fx.Node)
1087+
quant_input = get_arg(quant_node, "input", fx.Node)
1088+
mask_shape = get_shape(gm, quant_input)
1089+
if not mask_shape:
1090+
return None
1091+
mask_shape = list(mask_shape)
1092+
# Softmax mask is packed 16 elements per int32 word.
1093+
assert (
1094+
mask_shape[-1] % 16 == 0
1095+
), f"Softmax mask dimension must be divisible by 16, got {mask_shape[-1]}"
1096+
mask_shape[-1] = mask_shape[-1] // 16
1097+
mask_tensor = insert_node_with_meta(
1098+
gm,
1099+
torch.ops.aten.full.default,
1100+
(mask_shape, 0.0),
1101+
{"dtype": torch.int32},
1102+
anchor_node,
1103+
input_q,
1104+
)
1105+
# Initial position for streaming softmax (unused, set to 0).
1106+
pos_tensor = insert_node_with_meta(
1107+
gm,
1108+
torch.ops.aten.full.default,
1109+
([1], 0),
1110+
{"dtype": torch.int64},
1111+
anchor_node,
1112+
input_q,
1113+
)
1114+
args = (
1115+
input_q,
1116+
mask_tensor,
1117+
get_arg(anchor_node, "dim", int),
1118+
0,
1119+
pos_tensor,
1120+
get_arg(dq_input, "scale", float),
1121+
get_arg(dq_input, "zero_point", int),
1122+
get_arg(quant_node, "scale", float),
1123+
get_arg(quant_node, "zero_point", int),
1124+
)
1125+
return replace_with_op(
1126+
gm, anchor_node, self.replacement_op(), args, {}, quant_node
1127+
)
1128+
9881129

9891130
class MixedW8A32LinearPattern(QuantizationPattern):
9901131
def partition_types(self) -> List[OpOverload]:
@@ -1039,6 +1180,36 @@ def get_anchors(
10391180
def replacement_op(self) -> OpOverload:
10401181
return torch.ops.cadence.quantized_w8a32_linear.default
10411182

1183+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
1184+
if len(anchor_node.args) != 3 or len(anchor_node.kwargs) > 0:
1185+
return None
1186+
_arg1 = anchor_node.args[1]
1187+
dq_weight = (
1188+
_arg1
1189+
if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR
1190+
else None
1191+
)
1192+
_arg2 = anchor_node.args[2]
1193+
dq_bias = (
1194+
_arg2
1195+
if isinstance(_arg2, fx.Node) and _arg2.target == DQ_PER_TENSOR
1196+
else None
1197+
)
1198+
if dq_weight is None or dq_bias is None:
1199+
return None
1200+
input_node = anchor_node.args[0]
1201+
assert isinstance(input_node, fx.Node)
1202+
args = (
1203+
input_node,
1204+
get_arg(dq_weight, "input", fx.Node),
1205+
get_arg(dq_weight, "scale", float),
1206+
get_arg(dq_bias, "input", fx.Node),
1207+
get_arg(dq_bias, "scale", float),
1208+
)
1209+
return replace_with_op(
1210+
gm, anchor_node, self.replacement_op(), args, {}, anchor_node
1211+
)
1212+
10421213

10431214
class MixedW8A32ConvPattern(QuantizationPattern):
10441215
def partition_types(self) -> List[OpOverload]:
@@ -1113,6 +1284,57 @@ def get_anchors(
11131284
def replacement_op(self) -> OpOverload:
11141285
return torch.ops.cadence.quantized_w8a32_conv.default
11151286

1287+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
1288+
if len(anchor_node.args) != 3 or len(anchor_node.kwargs) > 0:
1289+
return None
1290+
_arg1 = anchor_node.args[1]
1291+
dq_weight = (
1292+
_arg1
1293+
if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR
1294+
else None
1295+
)
1296+
_arg2 = anchor_node.args[2]
1297+
dq_bias = (
1298+
_arg2
1299+
if isinstance(_arg2, fx.Node) and _arg2.target == DQ_PER_TENSOR
1300+
else None
1301+
)
1302+
if dq_weight is None or dq_bias is None:
1303+
return None
1304+
input_node = anchor_node.args[0]
1305+
assert isinstance(input_node, fx.Node)
1306+
assert get_arg(anchor_node, "stride", list[int]) == [1]
1307+
assert get_arg(anchor_node, "padding", list[int]) == [0]
1308+
assert get_arg(anchor_node, "dilation", list[int]) == [1]
1309+
assert get_arg(anchor_node, "groups", int) == 1
1310+
weight_q = get_arg(dq_weight, "input", fx.Node)
1311+
transposed_inputs = insert_node_with_meta(
1312+
gm,
1313+
torch.ops.aten.permute.default,
1314+
(input_node, [0, 2, 1]),
1315+
None,
1316+
anchor_node,
1317+
input_node,
1318+
)
1319+
transposed_weights = insert_node_with_meta(
1320+
gm,
1321+
torch.ops.aten.permute.default,
1322+
(weight_q, [2, 0, 1]),
1323+
None,
1324+
anchor_node,
1325+
weight_q,
1326+
)
1327+
args = (
1328+
transposed_inputs,
1329+
transposed_weights,
1330+
get_arg(dq_weight, "scale", float),
1331+
get_arg(dq_bias, "input", fx.Node),
1332+
get_arg(dq_bias, "scale", float),
1333+
)
1334+
return replace_with_op(
1335+
gm, anchor_node, self.replacement_op(), args, {}, anchor_node
1336+
)
1337+
11161338

11171339
class MixedW8A32GruPattern(QuantizationPattern):
11181340
def partition_types(self) -> List[OpOverload]:
@@ -1185,6 +1407,42 @@ def __init__(self, args, meta):
11851407
def replacement_op(self) -> OpOverload:
11861408
return torch.ops.cadence.quantized_w8a32_gru.default
11871409

1410+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None:
1411+
if len(anchor_node.kwargs) > 0:
1412+
return None
1413+
params = anchor_node.args[2]
1414+
# GRU requires 4 weight/bias params: w_ih, w_hh, b_ih, b_hh
1415+
if not isinstance(params, (list, tuple)) or len(params) < 4:
1416+
return None
1417+
dq_w_ih = params[0]
1418+
if not isinstance(dq_w_ih, fx.Node) or dq_w_ih.target != DQ_PER_TENSOR:
1419+
return None
1420+
dq_w_hh = params[1]
1421+
if not isinstance(dq_w_hh, fx.Node) or dq_w_hh.target != DQ_PER_TENSOR:
1422+
return None
1423+
dq_b_ih = params[2]
1424+
if not isinstance(dq_b_ih, fx.Node) or dq_b_ih.target != DQ_PER_TENSOR:
1425+
return None
1426+
dq_b_hh = params[3]
1427+
if not isinstance(dq_b_hh, fx.Node) or dq_b_hh.target != DQ_PER_TENSOR:
1428+
return None
1429+
input_node = anchor_node.args[0]
1430+
hidden_node = anchor_node.args[1]
1431+
args = (
1432+
input_node,
1433+
hidden_node,
1434+
get_arg(dq_w_ih, "input", fx.Node),
1435+
get_arg(dq_w_ih, "scale", float),
1436+
get_arg(dq_w_hh, "input", fx.Node),
1437+
get_arg(dq_w_hh, "scale", float),
1438+
get_arg(dq_b_ih, "input", fx.Node),
1439+
get_arg(dq_b_ih, "scale", float),
1440+
get_arg(dq_b_hh, "input", fx.Node),
1441+
)
1442+
return replace_with_op(
1443+
gm, anchor_node, self.replacement_op(), args, {}, anchor_node
1444+
)
1445+
11881446

11891447
class RmsNormPattern(QuantizationPattern):
11901448
"""Pattern that preserves rms_norm from decomposition without matching anything."""

0 commit comments

Comments
 (0)