Skip to content

Commit 0549e1e

Browse files
authored
Qualcomm AI Engine Direct - Keep a non-delegated op's weight dequantize on CPU (#23915)
### Summary FoldQDQ marks a quantized parameter's dequantize as a bypass node, so the partitioner force-passes it. When its consumer isn't delegated (rejected by op validation, or in `skip_node_id_set`), the dequantize is lowered alone as a QNN Dequantize, which rejects per-channel encodings, and context binary generation fails. The partitioner now leaves it on CPU with its consumer. Fixes #19908. ### Test plan `test_passes.py::TestPasses::test_skipped_conv_keeps_its_weight_dequantize_on_cpu` fails on main ("Op Dequantize does not support per-channel quant tensor") and passes here; the rest of `test_passes.py` passes on both. cc @cccclai @winskuo-quic @shewu-quic @haowhsu-quic @DannyYuyang-quic @cbilgin @abhinaykukkadapu @psiddh @nil-is-all
1 parent 8739d28 commit 0549e1e

2 files changed

Lines changed: 119 additions & 1 deletion

File tree

‎backends/qualcomm/partition/qnn_partitioner.py‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
PER_CHANNEL_GROUP_ENCODING,
1515
)
1616
from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
17+
from executorch.backends.qualcomm.builders.utils import is_parameter
1718
from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
1819
from executorch.backends.qualcomm.serialization.qc_schema import (
1920
HtpArch,
@@ -41,6 +42,7 @@
4142
PartitionResult,
4243
)
4344
from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
45+
from executorch.exir.dialects._ops import ops as exir_ops
4446
from torch.export.exported_program import ExportedProgram
4547
from torch.fx.passes.infra.partitioner import Partition
4648
from torch.fx.passes.operator_support import OperatorSupportBase
@@ -61,6 +63,13 @@
6163
logger = logging.getLogger(__name__)
6264
logger.setLevel(logging.DEBUG)
6365

66+
# dequantize ops with a CPU out variant; others cannot fall back
67+
CPU_DQ_OPS = {
68+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
69+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
70+
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
71+
}
72+
6473

6574
class QnnOperatorSupport(OperatorSupportBase):
6675
def __init__(
@@ -240,10 +249,42 @@ def generate_partitions(
240249
self.skip_node_id_set,
241250
self.skip_node_op_set,
242251
)
243-
return generate_partitions_from_list_of_nodes(
252+
partitions = generate_partitions_from_list_of_nodes(
244253
edge_program.graph_module,
245254
op_support=self.op_support_checker,
246255
)
256+
return self._drop_orphaned_parameter_dequants(partitions, edge_program)
257+
258+
@staticmethod
259+
def _drop_orphaned_parameter_dequants(
260+
partitions: List[Partition], edge_program: torch.export.ExportedProgram
261+
) -> List[Partition]:
262+
"""
263+
A parameter dequantize whose consumer is not delegated stays on CPU. A QNN
264+
Dequantize alone rejects per-channel encodings.
265+
"""
266+
for partition in partitions:
267+
stack = [
268+
node
269+
for node in partition.nodes
270+
if node.target in CPU_DQ_OPS
271+
and node.meta.get(QCOM_BYPASS_NODE, False)
272+
and is_parameter(node.args[0], edge_program)
273+
]
274+
while stack:
275+
node = stack.pop()
276+
if node not in partition.nodes or any(
277+
user in partition.nodes for user in node.users
278+
):
279+
continue
280+
partition.remove_node(node)
281+
stack.extend(
282+
arg
283+
for arg in node.args
284+
if isinstance(arg, torch.fx.Node)
285+
and arg.meta.get(QCOM_BYPASS_NODE, False)
286+
)
287+
return [partition for partition in partitions if partition.nodes]
247288

248289
def tag_nodes(
249290
self, partitions: List[Partition], edge_program: torch.export.ExportedProgram

‎backends/qualcomm/tests/test_passes.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -783,6 +783,83 @@ def forward(self, x):
783783
if "QNN" in str(e) or "qnn" in str(e):
784784
self.skipTest(f"QNN SDK not available: {e}")
785785

786+
def _assert_cpu_conv_keeps_its_weight_dequantize(
787+
self, quant_dtype, use_fp16, skip_node_id_set=None, emit=True
788+
):
789+
"""A quantized conv kept off QNN must take its per-channel weight dequantize
790+
with it. FoldQDQ force-passes that dequantize, so it used to be lowered alone
791+
as a QNN Dequantize, which rejects per-channel encodings and fails context
792+
binary generation (#19908)."""
793+
794+
class ConvDepthwiseConv(torch.nn.Module):
795+
def __init__(self):
796+
super().__init__()
797+
self.pw1 = torch.nn.Conv2d(8, 16, 1)
798+
self.dw = torch.nn.Conv2d(16, 16, 3, padding=1, groups=16)
799+
self.pw2 = torch.nn.Conv2d(16, 8, 1)
800+
801+
def forward(self, x):
802+
return self.pw2(torch.relu(self.dw(torch.relu(self.pw1(x)))))
803+
804+
sample_input = (torch.randn(1, 8, 16, 16),)
805+
quantizer = QnnQuantizer()
806+
quantizer.set_default_quant_config(
807+
quant_dtype=quant_dtype, is_conv_per_channel=True
808+
)
809+
prepared = prepare_pt2e(
810+
torch.export.export(ConvDepthwiseConv().eval(), sample_input).module(),
811+
quantizer,
812+
)
813+
prepared(*sample_input)
814+
compiler_specs = generate_qnn_executorch_compiler_spec(
815+
soc_model=QcomChipset.SM8650,
816+
backend_options=generate_htp_compiler_spec(use_fp16=use_fp16),
817+
)
818+
try:
819+
edge = to_edge_transform_and_lower_to_qnn(
820+
convert_pt2e(prepared),
821+
sample_input,
822+
compiler_specs,
823+
skip_node_id_set=skip_node_id_set,
824+
)
825+
except RuntimeError as e:
826+
if "QNN" in str(e) or "qnn" in str(e):
827+
self.skipTest(f"QNN SDK not available: {e}")
828+
raise
829+
830+
# inspected before to_executorch(), which swaps in out-variant ops
831+
def cpu_nodes(target):
832+
return [
833+
node
834+
for node in edge.exported_program().graph.nodes
835+
if node.op == "call_function" and node.target == target
836+
]
837+
838+
convs = cpu_nodes(exir_ops.edge.aten.convolution.default)
839+
self.assertEqual([conv.name for conv in convs], ["aten_convolution_default_1"])
840+
dequants = cpu_nodes(
841+
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default
842+
)
843+
self.assertTrue(dequants)
844+
for dequant in dequants:
845+
self.assertEqual(list(dequant.users), convs)
846+
if emit:
847+
edge.to_executorch()
848+
849+
def test_skipped_conv_keeps_its_weight_dequantize_on_cpu(self):
850+
self._assert_cpu_conv_keeps_its_weight_dequantize(
851+
QuantDtype.use_16a8w,
852+
use_fp16=False,
853+
skip_node_id_set={"aten_convolution_default_1"},
854+
)
855+
856+
def test_rejected_conv_keeps_its_weight_dequantize_on_cpu(self):
857+
# the HTP rejects fp16-activation/int8-weight depthwise conv. Not emitted:
858+
# InsertCastForFpActQuantizedWeight casts the CPU conv's bias meta to fp16.
859+
self._assert_cpu_conv_keeps_its_weight_dequantize(
860+
QuantDtype.use_fp16a8w, use_fp16=True, emit=False
861+
)
862+
786863
def test_lowering_accepts_exported_program_without_reexport(self):
787864
"""to_edge_transform_and_lower_to_qnn must lower a caller-supplied
788865
ExportedProgram as captured, never re-exporting it.

0 commit comments

Comments
 (0)