Skip to content

Commit 1f11fca

Browse files
DrJessopfacebook-github-bot
DrJessop
authored and
facebook-github-bot
committed
Duplicate cat if sandwiched between deq/quant
Summary: Suppose we have (int_data -> dq, fp32_data) -> cat -> (slice, q) This is a case with a multi-user cat. If we duplicated the cat, we would have (int_data -> dq, fp32_data) -> (cat -> slice, cat -> q) Then, in one of our chains, we could push the q above its cat. This allows for us to at least keep one chain using quantized math the whole way instead of forcing a break to fp32 in both chains. Example {F1990607611}. Just suppose that we already hoisted the quant node to be right underneath the cat. In a later pass, will hoist quant under single-user cat above the cat. This pass only duplicates the cat for each quant op user which has independent quant params. Reviewed By: mcremon-meta Differential Revision: D107174424
1 parent 40b0a35 commit 1f11fca

2 files changed

Lines changed: 331 additions & 0 deletions

File tree

‎backends/cadence/aot/reorder_ops.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -895,6 +895,79 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
895895
return should_swap(parent, node) and do_swap(parent, node)
896896

897897

898+
_QUANT_OVERLOAD_PACKETS = {
899+
exir_ops.edge.quantized_decomposed.quantize_per_tensor,
900+
exir_ops.edge.cadence.quantize_per_tensor,
901+
}
902+
903+
_DEQUANT_OVERLOAD_PACKETS = {
904+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor,
905+
exir_ops.edge.cadence.dequantize_per_tensor,
906+
}
907+
908+
909+
@register_cadence_pass(CadencePassAttribute(opt_level=1))
910+
class SplitDequantizedCatPass(RemoveOrReplacePassInterface):
911+
"""Split a cat node so that quantize consumers get their own copy.
912+
913+
Fires when a cat has all floating-point inputs, at least one dequantize
914+
input, and at least one quantize consumer. Quant consumers are grouped
915+
by matching qparams; each group receives a dedicated duplicate of the
916+
cat node. Non-quant consumers stay on the original cat, whose
917+
semantics are unchanged.
918+
919+
A later pass (e.g. AdvanceQuantizeOpAboveDefChainPass extended for cat)
920+
can then hoist each quant above its single-consumer cat copy without
921+
affecting the non-quant paths.
922+
"""
923+
924+
@property
925+
def targets(self) -> list[EdgeOpOverload]:
926+
return [exir_ops.edge.aten.cat.default]
927+
928+
def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
929+
cat_inputs = node.args[0]
930+
if not isinstance(cat_inputs, (list, tuple)):
931+
return False
932+
933+
has_dequant_input = False
934+
for inp in cat_inputs:
935+
assert isinstance(inp, torch.fx.Node)
936+
val = inp.meta["val"]
937+
if val is None or not val.is_floating_point():
938+
return False
939+
if get_overload_packet(inp.target) in _DEQUANT_OVERLOAD_PACKETS:
940+
has_dequant_input = True
941+
942+
if not has_dequant_input:
943+
return False
944+
945+
quant_groups: DefaultDict[Tuple, List[torch.fx.Node]] = defaultdict(list)
946+
for user in list(node.users.keys()):
947+
if get_overload_packet(user.target) in _QUANT_OVERLOAD_PACKETS:
948+
quant_groups[user.args[1:]].append(user)
949+
950+
if not quant_groups:
951+
return False
952+
953+
graph = node.graph
954+
dim = node.args[1] if len(node.args) > 1 else 0
955+
for quant_consumers in quant_groups.values():
956+
with graph.inserting_after(node):
957+
dup_cat = graph.call_function(
958+
node.target,
959+
args=(list(cat_inputs), dim),
960+
)
961+
# This pass will run export pass after to fix the type in
962+
# meta
963+
dup_cat.meta = node.meta.copy()
964+
965+
for q_node in quant_consumers:
966+
q_node.replace_input_with(node, dup_cat)
967+
968+
return True
969+
970+
898971
# The following class consolidates functions to reoder ops (i.e., either hoist
899972
# or sink some ops in the graph).
900973
class CadenceReorderOpsInGraph:

‎backends/cadence/aot/tests/test_reorder_ops_passes.py‎

Lines changed: 258 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
2929
PropagateSlice,
3030
SinkOpsCloserToUsePass,
31+
SplitDequantizedCatPass,
3132
)
3233
from executorch.backends.test.graph_builder import GraphBuilder
3334
from executorch.exir.dialects._ops import ops as exir_ops
@@ -1024,3 +1025,260 @@ def test_no_swap_binary_same_shape(self) -> None:
10241025
result = PropagateSlice().call(gm)
10251026

10261027
self.assertFalse(result.modified)
1028+
1029+
1030+
class TestSplitDequantizedCat(unittest.TestCase):
1031+
def test_no_dequant_input_noop(self) -> None:
1032+
"""Cat with only float (non-dequant) inputs should not be split."""
1033+
builder = GraphBuilder()
1034+
a = builder.placeholder("a", torch.randn(2, 4))
1035+
b = builder.placeholder("b", torch.randn(2, 4))
1036+
cat = builder.call_operator(
1037+
exir_ops.edge.aten.cat.default, args=([a, b], 0)
1038+
)
1039+
q = builder.call_operator(
1040+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1041+
args=(cat, 0.01, 0, -128, 127, torch.int8),
1042+
)
1043+
dq = builder.call_operator(
1044+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1045+
args=(q, 0.01, 0, -128, 127, torch.int8),
1046+
)
1047+
builder.output([dq])
1048+
gm = builder.get_graph_module()
1049+
1050+
result = SplitDequantizedCatPass().call(gm)
1051+
1052+
self.assertFalse(result.modified)
1053+
self.assertEqual(count_node(gm, exir_ops.edge.aten.cat.default), 1)
1054+
1055+
def test_no_quant_output_noop(self) -> None:
1056+
"""Cat with a dequant input but no quant consumer should not be split."""
1057+
builder = GraphBuilder()
1058+
x_int8 = builder.placeholder(
1059+
"x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
1060+
)
1061+
dq = builder.call_operator(
1062+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1063+
args=(x_int8, 0.01, 0, -128, 127, torch.int8),
1064+
)
1065+
b = builder.placeholder("b", torch.randn(2, 4))
1066+
cat = builder.call_operator(
1067+
exir_ops.edge.aten.cat.default, args=([dq, b], 0)
1068+
)
1069+
builder.output([cat])
1070+
gm = builder.get_graph_module()
1071+
1072+
result = SplitDequantizedCatPass().call(gm)
1073+
1074+
self.assertFalse(result.modified)
1075+
self.assertEqual(count_node(gm, exir_ops.edge.aten.cat.default), 1)
1076+
1077+
def test_one_dequant_input_one_quant_output(self) -> None:
1078+
"""Cat with one dequant input and one quant consumer should be split."""
1079+
builder = GraphBuilder()
1080+
x_int8 = builder.placeholder(
1081+
"x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
1082+
)
1083+
dq = builder.call_operator(
1084+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1085+
args=(x_int8, 0.01, 0, -128, 127, torch.int8),
1086+
)
1087+
b = builder.placeholder("b", torch.randn(2, 4))
1088+
cat = builder.call_operator(
1089+
exir_ops.edge.aten.cat.default, args=([dq, b], 0)
1090+
)
1091+
sliced = builder.call_operator(
1092+
exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
1093+
)
1094+
q = builder.call_operator(
1095+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1096+
args=(cat, 0.02, -5, -128, 127, torch.int8),
1097+
)
1098+
q_dq = builder.call_operator(
1099+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1100+
args=(q, 0.02, -5, -128, 127, torch.int8),
1101+
)
1102+
builder.output([sliced, q_dq])
1103+
gm = builder.get_graph_module()
1104+
1105+
result = cast(PassResult, SplitDequantizedCatPass().call(gm))
1106+
1107+
self.assertTrue(result.modified)
1108+
converted = result.graph_module
1109+
self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 2)
1110+
1111+
# The slice should still be on the original cat, which has no quant consumers.
1112+
slice_nodes = converted.graph.find_nodes(
1113+
op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
1114+
)
1115+
for node in slice_nodes:
1116+
cat_input = node.args[0]
1117+
self.assertEqual(cat_input.target, exir_ops.edge.aten.cat.default)
1118+
quant_users = [
1119+
u
1120+
for u in cat_input.users
1121+
if u.target
1122+
== exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
1123+
]
1124+
self.assertEqual(len(quant_users), 0)
1125+
1126+
def test_non_quant_consumers_stay_on_original_cat(self) -> None:
1127+
"""All non-quant consumers should remain on the original cat."""
1128+
builder = GraphBuilder()
1129+
x_int8 = builder.placeholder(
1130+
"x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
1131+
)
1132+
dq = builder.call_operator(
1133+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1134+
args=(x_int8, 0.01, 0, -128, 127, torch.int8),
1135+
)
1136+
b = builder.placeholder("b", torch.randn(2, 4))
1137+
cat = builder.call_operator(
1138+
exir_ops.edge.aten.cat.default, args=([dq, b], 0)
1139+
)
1140+
sliced = builder.call_operator(
1141+
exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
1142+
)
1143+
abs_val = builder.call_operator(
1144+
exir_ops.edge.aten.abs.default, args=(cat,)
1145+
)
1146+
q = builder.call_operator(
1147+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1148+
args=(cat, 0.02, -5, -128, 127, torch.int8),
1149+
)
1150+
q_dq = builder.call_operator(
1151+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1152+
args=(q, 0.02, -5, -128, 127, torch.int8),
1153+
)
1154+
builder.output([sliced, abs_val, q_dq])
1155+
gm = builder.get_graph_module()
1156+
1157+
result = cast(PassResult, SplitDequantizedCatPass().call(gm))
1158+
1159+
self.assertTrue(result.modified)
1160+
converted = result.graph_module
1161+
self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 2)
1162+
1163+
# Both non-quant consumers (slice and abs) should use the same cat.
1164+
slice_nodes = converted.graph.find_nodes(
1165+
op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
1166+
)
1167+
abs_nodes = converted.graph.find_nodes(
1168+
op="call_function", target=exir_ops.edge.aten.abs.default
1169+
)
1170+
self.assertEqual(len(slice_nodes), 1)
1171+
self.assertEqual(len(abs_nodes), 1)
1172+
self.assertIs(slice_nodes[0].args[0], abs_nodes[0].args[0])
1173+
1174+
# That shared cat should have no quant consumers.
1175+
original_cat = slice_nodes[0].args[0]
1176+
quant_users = [
1177+
u
1178+
for u in original_cat.users
1179+
if u.target
1180+
== exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
1181+
]
1182+
self.assertEqual(len(quant_users), 0)
1183+
1184+
def test_two_quant_outputs_same_params_shared_cat(self) -> None:
1185+
"""Two quant consumers with identical params should share one duplicate cat."""
1186+
builder = GraphBuilder()
1187+
x_int8 = builder.placeholder(
1188+
"x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
1189+
)
1190+
dq = builder.call_operator(
1191+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1192+
args=(x_int8, 0.01, 0, -128, 127, torch.int8),
1193+
)
1194+
b = builder.placeholder("b", torch.randn(2, 4))
1195+
cat = builder.call_operator(
1196+
exir_ops.edge.aten.cat.default, args=([dq, b], 0)
1197+
)
1198+
sliced = builder.call_operator(
1199+
exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
1200+
)
1201+
q1 = builder.call_operator(
1202+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1203+
args=(cat, 0.02, -5, -128, 127, torch.int8),
1204+
)
1205+
q2 = builder.call_operator(
1206+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1207+
args=(cat, 0.02, -5, -128, 127, torch.int8),
1208+
)
1209+
dq1 = builder.call_operator(
1210+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1211+
args=(q1, 0.02, -5, -128, 127, torch.int8),
1212+
)
1213+
dq2 = builder.call_operator(
1214+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1215+
args=(q2, 0.02, -5, -128, 127, torch.int8),
1216+
)
1217+
builder.output([sliced, dq1, dq2])
1218+
gm = builder.get_graph_module()
1219+
1220+
result = cast(PassResult, SplitDequantizedCatPass().call(gm))
1221+
1222+
self.assertTrue(result.modified)
1223+
converted = result.graph_module
1224+
# Original cat + one shared duplicate = 2 cats total
1225+
self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 2)
1226+
1227+
# Both quant nodes should share the same cat input (the duplicate).
1228+
quant_nodes = converted.graph.find_nodes(
1229+
op="call_function",
1230+
target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1231+
)
1232+
quant_cat_inputs = {node.args[0] for node in quant_nodes}
1233+
self.assertEqual(len(quant_cat_inputs), 1)
1234+
1235+
def test_two_quant_outputs_different_params_separate_cats(self) -> None:
1236+
"""Two quant consumers with different params should get separate duplicate cats."""
1237+
builder = GraphBuilder()
1238+
x_int8 = builder.placeholder(
1239+
"x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
1240+
)
1241+
dq = builder.call_operator(
1242+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1243+
args=(x_int8, 0.01, 0, -128, 127, torch.int8),
1244+
)
1245+
b = builder.placeholder("b", torch.randn(2, 4))
1246+
cat = builder.call_operator(
1247+
exir_ops.edge.aten.cat.default, args=([dq, b], 0)
1248+
)
1249+
sliced = builder.call_operator(
1250+
exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
1251+
)
1252+
q1 = builder.call_operator(
1253+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1254+
args=(cat, 0.02, -5, -128, 127, torch.int8),
1255+
)
1256+
q2 = builder.call_operator(
1257+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1258+
args=(cat, 0.03, 10, -128, 127, torch.int8),
1259+
)
1260+
dq1 = builder.call_operator(
1261+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1262+
args=(q1, 0.02, -5, -128, 127, torch.int8),
1263+
)
1264+
dq2 = builder.call_operator(
1265+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
1266+
args=(q2, 0.03, 10, -128, 127, torch.int8),
1267+
)
1268+
builder.output([sliced, dq1, dq2])
1269+
gm = builder.get_graph_module()
1270+
1271+
result = cast(PassResult, SplitDequantizedCatPass().call(gm))
1272+
1273+
self.assertTrue(result.modified)
1274+
converted = result.graph_module
1275+
# Original cat + two separate duplicates = 3 cats total
1276+
self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 3)
1277+
1278+
# Each quant node should have a different cat input.
1279+
quant_nodes = converted.graph.find_nodes(
1280+
op="call_function",
1281+
target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
1282+
)
1283+
quant_cat_inputs = {node.args[0] for node in quant_nodes}
1284+
self.assertEqual(len(quant_cat_inputs), 2)

0 commit comments

Comments
 (0)