Skip to content

Commit ed2d9ec

Browse files
winskuo-quicwinskuo-quic
authored
Qualcomm AI Engine Direct - Fix LPAI Partitioner QDQ Ops (pytorch#21838)
### Summary We made assumption that source node always saves quant_attrs target using quant_node, however, if a node's input is a constant, it will be using dequant_node's quant attrs. This assumption causes some error during node creation when are creating qdq nodes. To resolve this issue, we make quant and de-quant node explicit instead of relying on source node's quant attributes target. ### Test plan Passing UT for lpai paritioner
1 parent 601f1c2 commit ed2d9ec

4 files changed

Lines changed: 37 additions & 16 deletions

File tree

‎backends/qualcomm/_passes/insert_io_qdq.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
insert_quant_node,
1111
)
1212

13-
from executorch.backends.qualcomm.builders.node_visitor import q_dq_map, q_ops
13+
from executorch.backends.qualcomm.builders.node_visitor import q_ops, to_dq_op
1414

1515
from executorch.backends.qualcomm.builders.utils import (
1616
is_mutable_buffer_input,
@@ -76,7 +76,7 @@ def _insert(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
7676
graph_module=graph_module,
7777
input_node=n,
7878
output_node=user,
79-
target=q_dq_map[n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]],
79+
target=to_dq_op(n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]),
8080
)
8181

8282
def call(self, graph_module: torch.fx.GraphModule):

‎backends/qualcomm/_passes/lpai_partition_fallback_support.py‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,12 @@
1111
insert_dequant_node,
1212
insert_quant_node,
1313
)
14-
from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_dq_map, q_ops
14+
from executorch.backends.qualcomm.builders.node_visitor import (
15+
dq_ops,
16+
q_ops,
17+
to_dq_op,
18+
to_q_op,
19+
)
1520

1621
from executorch.backends.qualcomm.builders.utils import is_graph_input, is_graph_output
1722

@@ -108,15 +113,15 @@ def preserve_io_qdq(self, graph_module: torch.fx.GraphModule) -> None:
108113
graph_module=graph_module,
109114
input_node=n,
110115
output_node=user,
111-
target=n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING],
116+
target=to_q_op(n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]),
112117
pop_quant_attrs=False,
113118
)
114119
q_node.meta[QCOM_FALLBACK_NODE] = True
115120
dq_node = insert_dequant_node(
116121
graph_module=graph_module,
117122
input_node=q_node,
118123
output_node=user,
119-
target=q_dq_map[q_node.target],
124+
target=to_dq_op(q_node.target),
120125
)
121126
dq_node.meta[QCOM_BYPASS_NODE] = True
122127
elif (
@@ -131,15 +136,17 @@ def preserve_io_qdq(self, graph_module: torch.fx.GraphModule) -> None:
131136
graph_module=graph_module,
132137
input_node=output_node,
133138
output_node=getitem_node,
134-
target=output_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING],
139+
target=to_q_op(
140+
output_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]
141+
),
135142
pop_quant_attrs=False,
136143
)
137144
q_node.meta[QCOM_BYPASS_NODE] = True
138145
dq_node = insert_dequant_node(
139146
graph_module=graph_module,
140147
input_node=q_node,
141148
output_node=getitem_node,
142-
target=q_dq_map[q_node.target],
149+
target=to_dq_op(q_node.target),
143150
)
144151
dq_node.meta[QCOM_FALLBACK_NODE] = True
145152

@@ -191,7 +198,7 @@ def insert_partition_qdq(
191198
graph_module=graph_module,
192199
input_node=input_node,
193200
output_node=node,
194-
target=input_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING],
201+
target=to_q_op(input_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]),
195202
pop_quant_attrs=False,
196203
)
197204
for input_node in input_nodes
@@ -203,7 +210,7 @@ def insert_partition_qdq(
203210
graph_module=graph_module,
204211
input_node=input_q_node,
205212
output_node=node,
206-
target=q_dq_map[input_q_node.target],
213+
target=to_dq_op(input_q_node.target),
207214
)
208215
for input_q_node in input_q_nodes
209216
]
@@ -232,17 +239,17 @@ def insert_partition_qdq(
232239
graph_module=graph_module,
233240
input_node=output_node,
234241
output_node=output_user_node,
235-
target=output_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING],
242+
target=to_q_op(output_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]),
236243
pop_quant_attrs=False,
237244
)
238245
output_q_node.meta[QCOM_FALLBACK_NODE] = True
239246
output_dq_node = insert_dequant_node(
240247
graph_module=graph_module,
241248
input_node=output_q_node,
242249
output_node=output_user_node,
243-
target=q_dq_map[
250+
target=to_dq_op(
244251
output_q_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]
245-
],
252+
),
246253
)
247254
output_dq_node.meta[QCOM_BYPASS_NODE] = True
248255
graph_module.graph.eliminate_dead_code()

‎backends/qualcomm/_passes/utils.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from typing import Callable, Dict, List
88

99
import torch
10-
from executorch.backends.qualcomm.builders.node_visitor import q_ops
10+
from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops
1111
from executorch.backends.qualcomm.builders.utils import get_parameter
1212
from executorch.backends.qualcomm.utils.constants import (
1313
QCOM_DTYPE,
@@ -81,6 +81,9 @@ def insert_quant_node(
8181
quant_attrs: Dict = None,
8282
pop_quant_attrs: bool = True,
8383
) -> torch.fx.Node:
84+
assert (
85+
target in q_ops
86+
), f"insert_quant_node expects a quantize target, got: {target}"
8487
with graph_module.graph.inserting_after(input_node):
8588
inserted_node = _create_q_or_dq_node(
8689
graph_module=graph_module,
@@ -101,6 +104,9 @@ def insert_dequant_node(
101104
output_node: torch.fx.node,
102105
target: torch.fx.node.Target,
103106
) -> None:
107+
assert (
108+
target in dq_ops
109+
), f"insert_dequant_node expects a dequantize target, got: {target}"
104110
with graph_module.graph.inserting_after(input_node):
105111
inserted_node = _create_q_or_dq_node(
106112
graph_module=graph_module, node=input_node, target=target

‎backends/qualcomm/builders/node_visitor.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -101,13 +101,21 @@
101101
q_dq_map = {
102102
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
103103
exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
104-
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
105-
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
106104
exir_ops.edge.quantized_decomposed.quantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
107-
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
105+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
106+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor: exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor,
107+
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default: exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
108108
}
109109

110110

111+
def to_q_op(target):
112+
return target if target in q_ops else q_dq_map[target]
113+
114+
115+
def to_dq_op(target):
116+
return target if target in dq_ops else q_dq_map[target]
117+
118+
111119
class NodeVisitor:
112120
"""
113121
Node visitor pattern for visiting nodes in an edge IR graph

0 commit comments

Comments
 (0)