Skip to content

Commit be8df33

Browse files
rascanirascani
authored
Transforms: keep permutes around per-channel QDQ (pytorch#21894)
### Summary Stop FuseTransposeOrPermuteOpPairsPass from bypassing per-channel quantize and dequantize operators. Removing a surrounding permutation changes which logical dimension the stored axis selects, and this pass does not rewrite the axis argument. Per-tensor QDQ remains bypassable because its scalar parameters are layout-invariant. The per-channel entries originated in the Cadence-local transpose pass without axis-remapping logic or dedicated coverage, then moved mechanically into the shared transform. Keep the common bypass contract limited to layout-invariant operators; recovering axis-stable or explicitly remapped per-channel cases can be handled by a future axis-aware optimization. Add regressions for positive per-tensor fusion, individual and chained per-channel QDQ, mixed per-tensor/per-channel forks, and Cadence inheritance. These tests preserve the useful layout-invariant optimization while proving that an axis-bearing branch prevents unsafe producer removal. ### Test Plan Validated pytest -q backends/transforms/test/test_permute_optimization_passes.py: 51 passed. The Cadence regression was statically reviewed but could not be collected locally because the available Python 3.11 environment lacks typing.override; run it under Python 3.12 or the Cadence Buck target. Authored with Codex.
1 parent 9ffbe36 commit be8df33

4 files changed

Lines changed: 237 additions & 4 deletions

File tree

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -947,6 +947,51 @@ def _create_operator(
947947
else:
948948
raise ValueError(f"Unsupported op: {op}")
949949

950+
def test_per_channel_qdq_is_not_bypassed_without_axis_remap(self) -> None:
951+
for op, x_data in (
952+
(
953+
exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
954+
torch.randn(1, 2, 3, 4),
955+
),
956+
(
957+
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
958+
torch.randint(-128, 127, (1, 2, 3, 4), dtype=torch.int8),
959+
),
960+
):
961+
with self.subTest(op=op):
962+
builder = GraphBuilder()
963+
x = builder.placeholder("x", x_data)
964+
scales = builder.placeholder("scales", torch.tensor([0.25, 0.5]))
965+
zero_points = builder.placeholder(
966+
"zero_points", torch.tensor([0, 0], dtype=torch.int64)
967+
)
968+
to_nhwc = builder.call_operator(
969+
op=exir_ops.edge.aten.permute_copy.default,
970+
args=(x, [0, 2, 3, 1]),
971+
)
972+
qdq = builder.call_operator(
973+
op=op,
974+
args=(to_nhwc, scales, zero_points, 3, -128, 127, torch.int8),
975+
)
976+
to_nchw = builder.call_operator(
977+
op=exir_ops.edge.aten.permute_copy.default,
978+
args=(qdq, [0, 3, 1, 2]),
979+
)
980+
builder.output([to_nchw])
981+
graph_module = builder.get_graph_module()
982+
983+
result = cast(
984+
PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module)
985+
)
986+
987+
self.assertFalse(result.modified)
988+
self.assertEqual(
989+
count_node(
990+
result.graph_module, exir_ops.edge.aten.permute_copy.default
991+
),
992+
2,
993+
)
994+
950995

951996
class TestFuseTransposeOpPairsPass(TestFusionPassesBase):
952997
def _create_operator(

‎backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,12 @@ class FuseTransposeOrPermuteOpPairsPass(FuseOpPairsAcrossBranchesPass):
2929
so transpose(1, 2) then transpose(0, 2) is a pseudo identity and should be fused.
3030
"""
3131

32-
# A list of ops that can be bypassed when looking for a
33-
# transpose-permute chain. Subclasses can extend this with backend-specific ops.
32+
# Only layout-invariant ops can be bypassed. Per-channel QDQ requires
33+
# remapping its axis and must not be added here without that rewrite.
34+
# Subclasses can extend this with backend-specific ops.
3435
bypass_ops: set[EdgeOpOverload] = {
3536
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
36-
exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
3737
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
38-
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
3938
}
4039

4140
def can_fuse_for_chain(

‎backends/transforms/targets.bzl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -473,6 +473,7 @@ def define_common_targets():
473473
"//executorch/exir/dialects:lib",
474474
":fuse_cascaded_transpose_or_permute_ops",
475475
":fuse_cascaded_view_ops",
476+
":fuse_transpose_or_permute_op_pairs_pass",
476477
":postpone_permute_below_squeeze_view",
477478
":remove_permutes_around_elementwise_ops",
478479
":replace_nop_transpose_or_permute_with_view",

‎backends/transforms/test/test_permute_optimization_passes.py‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@
1717
FuseCascadedTransposeOrPermuteOps,
1818
)
1919
from executorch.backends.transforms.fuse_cascaded_view_ops import FuseCascadedViewOps
20+
from executorch.backends.transforms.fuse_transpose_or_permute_op_pairs_pass import (
21+
FuseTransposeOrPermuteOpPairsPass,
22+
)
2023
from executorch.backends.transforms.postpone_permute_below_squeeze_view import (
2124
PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
2225
)
@@ -551,6 +554,191 @@ def test_negative_not_squeeze_like(self) -> None:
551554
self.assertEqual(targets[1], exir_ops.edge.aten.permute_copy.default)
552555

553556

557+
class FuseTransposeOrPermuteOpPairsTest(unittest.TestCase):
558+
def test_per_tensor_qdq_is_bypassed(self) -> None:
559+
for op, x_data in (
560+
(
561+
exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
562+
torch.randn(1, 2, 3, 4),
563+
),
564+
(
565+
exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
566+
torch.randint(-128, 127, (1, 2, 3, 4), dtype=torch.int8),
567+
),
568+
):
569+
with self.subTest(op=op):
570+
builder = GraphBuilder()
571+
x = builder.placeholder("x", x_data)
572+
to_nhwc = builder.call_operator(
573+
op=exir_ops.edge.aten.permute_copy.default,
574+
args=(x, [0, 2, 3, 1]),
575+
)
576+
qdq = builder.call_operator(
577+
op=op,
578+
args=(to_nhwc, 0.25, 0, -128, 127, torch.int8),
579+
)
580+
to_nchw = builder.call_operator(
581+
op=exir_ops.edge.aten.permute_copy.default,
582+
args=(qdq, [0, 3, 1, 2]),
583+
)
584+
builder.output([to_nchw])
585+
graph_module = builder.get_graph_module()
586+
before = copy.deepcopy(graph_module)
587+
588+
result = cast(
589+
PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module)
590+
)
591+
592+
self.assertTrue(result.modified)
593+
self.assertEqual(
594+
count_node(
595+
result.graph_module, exir_ops.edge.aten.permute_copy.default
596+
),
597+
0,
598+
)
599+
validate_numerics(
600+
before,
601+
result.graph_module,
602+
[x_data],
603+
"FuseTransposeOrPermuteOpPairsPass",
604+
)
605+
606+
def test_per_channel_qdq_is_not_bypassed_without_axis_remap(self) -> None:
607+
for op, x_data in (
608+
(
609+
exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
610+
torch.randn(1, 2, 3, 4),
611+
),
612+
(
613+
exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
614+
torch.randint(-128, 127, (1, 2, 3, 4), dtype=torch.int8),
615+
),
616+
):
617+
with self.subTest(op=op):
618+
builder = GraphBuilder()
619+
x = builder.placeholder("x", x_data)
620+
scales = builder.placeholder("scales", torch.tensor([0.25, 0.5]))
621+
zero_points = builder.placeholder(
622+
"zero_points", torch.tensor([0, 0], dtype=torch.int64)
623+
)
624+
to_nhwc = builder.call_operator(
625+
op=exir_ops.edge.aten.permute_copy.default,
626+
args=(x, [0, 2, 3, 1]),
627+
)
628+
qdq = builder.call_operator(
629+
op=op,
630+
# NHWC axis 3 would need to become NCHW axis 1 if the
631+
# surrounding permutes were removed.
632+
args=(to_nhwc, scales, zero_points, 3, -128, 127, torch.int8),
633+
)
634+
to_nchw = builder.call_operator(
635+
op=exir_ops.edge.aten.permute_copy.default,
636+
args=(qdq, [0, 3, 1, 2]),
637+
)
638+
builder.output([to_nchw])
639+
graph_module = builder.get_graph_module()
640+
before = copy.deepcopy(graph_module)
641+
642+
result = cast(
643+
PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module)
644+
)
645+
646+
self.assertFalse(result.modified)
647+
self.assertEqual(
648+
count_node(
649+
result.graph_module, exir_ops.edge.aten.permute_copy.default
650+
),
651+
2,
652+
)
653+
validate_numerics(
654+
before,
655+
result.graph_module,
656+
[x_data, torch.tensor([0.25, 0.5]), torch.tensor([0, 0])],
657+
"FuseTransposeOrPermuteOpPairsPass",
658+
)
659+
660+
def test_per_channel_qdq_chain_is_not_bypassed(self) -> None:
661+
builder = GraphBuilder()
662+
x_data = torch.randn(1, 2, 3, 4)
663+
scales_data = torch.tensor([0.25, 0.5])
664+
zero_points_data = torch.tensor([0, 0], dtype=torch.int64)
665+
x = builder.placeholder("x", x_data)
666+
scales = builder.placeholder("scales", scales_data)
667+
zero_points = builder.placeholder("zero_points", zero_points_data)
668+
to_nhwc = builder.call_operator(
669+
op=exir_ops.edge.aten.permute_copy.default,
670+
args=(x, [0, 2, 3, 1]),
671+
)
672+
quantize = builder.call_operator(
673+
op=exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
674+
args=(to_nhwc, scales, zero_points, 3, -128, 127, torch.int8),
675+
)
676+
dequantize = builder.call_operator(
677+
op=exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
678+
args=(quantize, scales, zero_points, 3, -128, 127, torch.int8),
679+
)
680+
to_nchw = builder.call_operator(
681+
op=exir_ops.edge.aten.permute_copy.default,
682+
args=(dequantize, [0, 3, 1, 2]),
683+
)
684+
builder.output([to_nchw])
685+
graph_module = builder.get_graph_module()
686+
before = copy.deepcopy(graph_module)
687+
688+
result = cast(PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module))
689+
690+
self.assertFalse(result.modified)
691+
self.assertEqual(
692+
count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default),
693+
2,
694+
)
695+
validate_numerics(
696+
before,
697+
result.graph_module,
698+
[x_data, scales_data, zero_points_data],
699+
"FuseTransposeOrPermuteOpPairsPass",
700+
)
701+
702+
def test_per_channel_branch_blocks_shared_permute_fusion(self) -> None:
703+
builder = GraphBuilder()
704+
x_data = torch.randn(1, 2, 3, 4)
705+
scales_data = torch.tensor([0.25, 0.5])
706+
zero_points_data = torch.tensor([0, 0], dtype=torch.int64)
707+
x = builder.placeholder("x", x_data)
708+
scales = builder.placeholder("scales", scales_data)
709+
zero_points = builder.placeholder("zero_points", zero_points_data)
710+
to_nhwc = builder.call_operator(
711+
op=exir_ops.edge.aten.permute_copy.default,
712+
args=(x, [0, 2, 3, 1]),
713+
)
714+
per_tensor = builder.call_operator(
715+
op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
716+
args=(to_nhwc, 0.25, 0, -128, 127, torch.int8),
717+
)
718+
per_channel = builder.call_operator(
719+
op=exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
720+
args=(to_nhwc, scales, zero_points, 3, -128, 127, torch.int8),
721+
)
722+
per_tensor_out = builder.call_operator(
723+
op=exir_ops.edge.aten.permute_copy.default,
724+
args=(per_tensor, [0, 3, 1, 2]),
725+
)
726+
per_channel_out = builder.call_operator(
727+
op=exir_ops.edge.aten.permute_copy.default,
728+
args=(per_channel, [0, 3, 1, 2]),
729+
)
730+
builder.output([per_tensor_out, per_channel_out])
731+
graph_module = builder.get_graph_module()
732+
733+
result = cast(PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module))
734+
735+
self.assertFalse(result.modified)
736+
self.assertEqual(
737+
count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default),
738+
3,
739+
)
740+
741+
554742
# ──────────────────────────────────────────────────────────────────────
555743
# Tests for ReplaceNopTransposeOrPermuteWithViewPass
556744
# ──────────────────────────────────────────────────────────────────────

0 commit comments

Comments
 (0)