Skip to content

Commit e65a3a3

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 Differential Revision: D105728177
1 parent db0a643 commit e65a3a3

1 file changed

Lines changed: 257 additions & 3 deletions

File tree

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

Lines changed: 257 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from typing import List, Optional, 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
1617
from executorch.backends.cadence.aot.quantizer.pattern_utils import (
1718
find_quant_user,
@@ -25,6 +26,7 @@
2526
check_out_zero_point_is_min_range,
2627
copy_node_metadata,
2728
get_bias_qparams,
29+
quantize_tensor_multiplier,
2830
)
2931
from torch import fx
3032
from torch._ops import OpOverload
@@ -304,9 +306,8 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
304306
quant_node = find_quant_user(relu_node)
305307
if quant_node is None:
306308
return None
307-
check_out_zero_point_is_min_range(
308-
quant_node.args[2], quant_node.args[5] # pyre-ignore[6]
309-
)
309+
# pyre-ignore[6]: Argument -> int/dtype narrowing
310+
check_out_zero_point_is_min_range(quant_node.args[2], quant_node.args[5])
310311
args = (
311312
dq0.args[0],
312313
dq0.args[1],
@@ -803,6 +804,30 @@ def get_anchors(
803804
def replacement_op(self) -> OpOverload:
804805
return torch.ops.cadence.quantized_max_pool2d_nchw.default
805806

807+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
808+
dq_input = get_dequant(anchor_node.args[0])
809+
if dq_input is None:
810+
return None
811+
quant_node = find_quant_user(anchor_node)
812+
if quant_node is None:
813+
return None
814+
kernel_size = get_arg(anchor_node, "kernel_size", Optional[list[int]]) or [1, 1]
815+
stride = get_arg(anchor_node, "stride", Optional[list[int]]) or kernel_size
816+
padding = get_arg(anchor_node, "padding", Optional[list[int]]) or [0, 0]
817+
dilation = get_arg(anchor_node, "dilation", Optional[list[int]]) or [1, 1]
818+
ceil_mode = get_arg(anchor_node, "ceil_mode", Optional[bool]) or False
819+
args = (dq_input.args[0],)
820+
kwargs = {
821+
"kernel_size": kernel_size,
822+
"stride": stride,
823+
"padding": padding,
824+
"dilation": dilation,
825+
"ceil_mode": ceil_mode,
826+
}
827+
return insert_fused_op(
828+
gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
829+
)
830+
806831

807832
class MaxPool2dWithoutIndicesPattern(QuantizationPattern):
808833
"""
@@ -842,6 +867,30 @@ def get_anchors(
842867
def replacement_op(self) -> OpOverload:
843868
return torch.ops.cadence.quantized_max_pool2d_nchw.default
844869

870+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
871+
dq_input = get_dequant(anchor_node.args[0])
872+
if dq_input is None:
873+
return None
874+
quant_node = find_quant_user(anchor_node)
875+
if quant_node is None:
876+
return None
877+
kernel_size = get_arg(anchor_node, "kernel_size", Optional[list[int]]) or [1, 1]
878+
stride = get_arg(anchor_node, "stride", Optional[list[int]]) or kernel_size
879+
padding = get_arg(anchor_node, "padding", Optional[list[int]]) or [0, 0]
880+
dilation = get_arg(anchor_node, "dilation", Optional[list[int]]) or [1, 1]
881+
ceil_mode = get_arg(anchor_node, "ceil_mode", Optional[bool]) or False
882+
args = (dq_input.args[0],)
883+
kwargs = {
884+
"kernel_size": kernel_size,
885+
"stride": stride,
886+
"padding": padding,
887+
"dilation": dilation,
888+
"ceil_mode": ceil_mode,
889+
}
890+
return insert_fused_op(
891+
gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
892+
)
893+
845894

846895
# This is a base class for ReLU
847896

@@ -871,6 +920,29 @@ def get_anchors(
871920
def replacement_op(self) -> OpOverload:
872921
return torch.ops.cadence.quantized_relu.per_tensor
873922

923+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
924+
dq_input = get_dequant(anchor_node.args[0])
925+
if dq_input is None:
926+
return None
927+
quant_node = find_quant_user(anchor_node)
928+
if quant_node is None:
929+
return None
930+
input_scale = dq_input.args[1]
931+
# pyre-fixme[58]
932+
requantize_scale = input_scale / quant_node.args[1]
933+
requantize_scale_t = torch.tensor([requantize_scale])
934+
out_multiplier, out_shift = quantize_tensor_multiplier(requantize_scale_t)
935+
args = (dq_input.args[0],)
936+
kwargs = {
937+
"X_zero_point": dq_input.args[2],
938+
"out_zero_point": quant_node.args[2],
939+
"out_multiplier": out_multiplier[0].item(),
940+
"out_shift": out_shift[0].item(),
941+
}
942+
return insert_fused_op(
943+
gm, anchor_node, self.replacement_op(), args, kwargs, quant_node
944+
)
945+
874946

875947
# Regular relu op
876948
class ReluPattern0(ReluBasePattern):
@@ -930,6 +1002,27 @@ def get_anchors(
9301002
def replacement_op(self) -> OpOverload:
9311003
return torch.ops.cadence.quantized_conv2d_nchw.per_tensor
9321004

1005+
def anchor_ops(self) -> tuple[OpOverload, ...]:
1006+
return (self.partition_types()[0],)
1007+
1008+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
1009+
conv_users = list(anchor_node.users)
1010+
if len(conv_users) != 1:
1011+
return None
1012+
relu_node = conv_users[0]
1013+
if relu_node.target != self.partition_types()[1]:
1014+
return None
1015+
dq_input = get_dequant(anchor_node.args[0])
1016+
dq_weight = get_dequant(anchor_node.args[1])
1017+
if dq_input is None or dq_weight is None:
1018+
return None
1019+
quant_node = find_quant_user(relu_node)
1020+
if quant_node is None:
1021+
return None
1022+
# pyre-ignore[6]: Argument -> int/dtype narrowing
1023+
check_out_zero_point_is_min_range(quant_node.args[2], quant_node.args[5])
1024+
return fuse_conv(self, gm, anchor_node, dq_input, dq_weight, quant_node)
1025+
9331026

9341027
# Conv1d + regular relu op fusion
9351028
class Conv1dReluPattern0(ConvReluBasePattern):
@@ -984,6 +1077,52 @@ def get_anchors(
9841077
def replacement_op(self) -> OpOverload:
9851078
return torch.ops.cadence.quantized_softmax.per_tensor
9861079

1080+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
1081+
dq_input = get_dequant(anchor_node.args[0])
1082+
if dq_input is None:
1083+
return None
1084+
quant_node = find_quant_user(anchor_node)
1085+
if quant_node is None:
1086+
return None
1087+
input_q = dq_input.args[0]
1088+
assert isinstance(input_q, fx.Node)
1089+
quant_input = quant_node.args[0]
1090+
assert isinstance(quant_input, fx.Node)
1091+
mask_shape = get_shape(gm, quant_input)
1092+
mask_shape = list(mask_shape) if mask_shape else []
1093+
mask_shape[-1] = mask_shape[-1] // 16
1094+
with gm.graph.inserting_before(anchor_node):
1095+
mask_tensor = gm.graph.call_function(
1096+
torch.ops.aten.full.default, (mask_shape, 0.0), {"dtype": torch.int32}
1097+
)
1098+
assert "val" in input_q.meta
1099+
fake_mode = input_q.meta["val"].fake_mode
1100+
assert fake_mode is not None
1101+
with fake_mode:
1102+
mask_tensor.meta["val"] = torch.full(mask_shape, 0.0, dtype=torch.int32)
1103+
copy_node_metadata(mask_tensor, input_q)
1104+
with gm.graph.inserting_before(anchor_node):
1105+
pos_tensor = gm.graph.call_function(
1106+
torch.ops.aten.full.default, ([1], 0), {"dtype": torch.int64}
1107+
)
1108+
with fake_mode:
1109+
pos_tensor.meta["val"] = torch.full([1], 0, dtype=torch.int64)
1110+
copy_node_metadata(pos_tensor, input_q)
1111+
args = (
1112+
input_q,
1113+
mask_tensor,
1114+
get_arg(anchor_node, "dim", int),
1115+
0,
1116+
pos_tensor,
1117+
dq_input.args[1],
1118+
dq_input.args[2],
1119+
quant_node.args[1],
1120+
quant_node.args[2],
1121+
)
1122+
return insert_fused_op(
1123+
gm, anchor_node, self.replacement_op(), args, {}, quant_node
1124+
)
1125+
9871126

9881127
class MixedW8A32LinearPattern(QuantizationPattern):
9891128
def partition_types(self) -> List[OpOverload]:
@@ -1038,6 +1177,28 @@ def get_anchors(
10381177
def replacement_op(self) -> OpOverload:
10391178
return torch.ops.cadence.quantized_w8a32_linear.default
10401179

1180+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
1181+
if len(anchor_node.args) != 3 or len(anchor_node.kwargs) > 0:
1182+
return None
1183+
dq_weight = get_dequant(anchor_node.args[1])
1184+
dq_bias = get_dequant(anchor_node.args[2])
1185+
if dq_weight is None or dq_bias is None:
1186+
return None
1187+
input_node = anchor_node.args[0]
1188+
assert isinstance(input_node, fx.Node)
1189+
args = (
1190+
input_node,
1191+
dq_weight.args[0],
1192+
dq_weight.args[1],
1193+
dq_bias.args[0],
1194+
dq_bias.args[1],
1195+
)
1196+
with gm.graph.inserting_after(anchor_node):
1197+
fused = gm.graph.call_function(self.replacement_op(), args, {})
1198+
fused.meta = anchor_node.meta
1199+
anchor_node.replace_all_uses_with(fused)
1200+
return fused
1201+
10411202

10421203
class MixedW8A32ConvPattern(QuantizationPattern):
10431204
def partition_types(self) -> List[OpOverload]:
@@ -1112,6 +1273,68 @@ def get_anchors(
11121273
def replacement_op(self) -> OpOverload:
11131274
return torch.ops.cadence.quantized_w8a32_conv.default
11141275

1276+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
1277+
if len(anchor_node.args) != 3 or len(anchor_node.kwargs) > 0:
1278+
return None
1279+
dq_weight = get_dequant(anchor_node.args[1])
1280+
dq_bias = get_dequant(anchor_node.args[2])
1281+
if dq_weight is None or dq_bias is None:
1282+
return None
1283+
input_node = anchor_node.args[0]
1284+
assert isinstance(input_node, fx.Node)
1285+
assert get_arg(anchor_node, "stride", list[int]) == [1]
1286+
assert get_arg(anchor_node, "padding", list[int]) == [0]
1287+
assert get_arg(anchor_node, "dilation", list[int]) == [1]
1288+
assert get_arg(anchor_node, "groups", int) == 1
1289+
weight_q = dq_weight.args[0]
1290+
assert isinstance(weight_q, fx.Node)
1291+
with gm.graph.inserting_before(anchor_node):
1292+
transposed_inputs = gm.graph.call_function(
1293+
torch.ops.aten.permute.default, (input_node, [0, 2, 1])
1294+
)
1295+
if "val" in input_node.meta:
1296+
original_val = input_node.meta["val"]
1297+
fake_mode = original_val.fake_mode
1298+
if fake_mode is not None:
1299+
with fake_mode:
1300+
transposed_inputs.meta["val"] = torch.ops.aten.permute.default(
1301+
original_val, [0, 2, 1]
1302+
)
1303+
else:
1304+
transposed_inputs.meta["val"] = torch.ops.aten.permute.default(
1305+
original_val, [0, 2, 1]
1306+
)
1307+
copy_node_metadata(transposed_inputs, input_node)
1308+
with gm.graph.inserting_before(anchor_node):
1309+
transposed_weights = gm.graph.call_function(
1310+
torch.ops.aten.permute.default, (weight_q, [2, 0, 1])
1311+
)
1312+
if "val" in weight_q.meta:
1313+
original_val = weight_q.meta["val"]
1314+
fake_mode = original_val.fake_mode
1315+
if fake_mode is not None:
1316+
with fake_mode:
1317+
transposed_weights.meta["val"] = torch.ops.aten.permute.default(
1318+
original_val, [2, 0, 1]
1319+
)
1320+
else:
1321+
transposed_weights.meta["val"] = torch.ops.aten.permute.default(
1322+
original_val, [2, 0, 1]
1323+
)
1324+
copy_node_metadata(transposed_weights, weight_q)
1325+
args = (
1326+
transposed_inputs,
1327+
transposed_weights,
1328+
dq_weight.args[1],
1329+
dq_bias.args[0],
1330+
dq_bias.args[1],
1331+
)
1332+
with gm.graph.inserting_after(anchor_node):
1333+
fused = gm.graph.call_function(self.replacement_op(), args, {})
1334+
fused.meta = anchor_node.meta
1335+
anchor_node.replace_all_uses_with(fused)
1336+
return fused
1337+
11151338

11161339
class MixedW8A32GruPattern(QuantizationPattern):
11171340
def partition_types(self) -> List[OpOverload]:
@@ -1184,6 +1407,37 @@ def __init__(self, args, meta):
11841407
def replacement_op(self) -> OpOverload:
11851408
return torch.ops.cadence.quantized_w8a32_gru.default
11861409

1410+
def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> Optional[fx.Node]:
1411+
if len(anchor_node.kwargs) > 0:
1412+
return None
1413+
params = anchor_node.args[2]
1414+
if not isinstance(params, (list, tuple)) or len(params) < 4:
1415+
return None
1416+
dq_w_ih = get_dequant(params[0])
1417+
dq_w_hh = get_dequant(params[1])
1418+
dq_b_ih = get_dequant(params[2])
1419+
dq_b_hh = get_dequant(params[3])
1420+
if dq_w_ih is None or dq_w_hh is None or dq_b_ih is None or dq_b_hh is None:
1421+
return None
1422+
input_node = anchor_node.args[0]
1423+
hidden_node = anchor_node.args[1]
1424+
args = (
1425+
input_node,
1426+
hidden_node,
1427+
dq_w_ih.args[0],
1428+
dq_w_ih.args[1],
1429+
dq_w_hh.args[0],
1430+
dq_w_hh.args[1],
1431+
dq_b_ih.args[0],
1432+
dq_b_ih.args[1],
1433+
dq_b_hh.args[0],
1434+
)
1435+
with gm.graph.inserting_after(anchor_node):
1436+
fused = gm.graph.call_function(self.replacement_op(), args, {})
1437+
fused.meta = anchor_node.meta
1438+
anchor_node.replace_all_uses_with(fused)
1439+
return fused
1440+
11871441

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

0 commit comments

Comments
 (0)