Skip to content

Commit bc89f05

Browse files
Add transpose variants to XNNPACK shared-qparam op list
aten.transpose.int and its siblings (transpose_copy.int, t.default, t_copy.default, swapaxes.default) were missing from _is_share_obs_or_fq_op in xnnpack_quantizer_utils.py, even though the value-preserving sibling permute/permute_copy was already present. Without this fix, propagate_annotation does not attach a SharedQuantizationSpec to transpose, so the input and output activations can be observed independently and receive different (scale, zero_point). After to_edge decomposes aten.transpose.int to aten.permute_copy.default, XNNPACK lowers it to XNNStaticTranspose, which hard-requires identical input/output quantization parameters for qint8/quint8 and rejects a mismatch with xnn_status_invalid_parameter. ARM (TOSA) and QNN quantizers already treat transpose as a shared-qparam op; this aligns the XNNPACK quantizer with them. Test: test_transpose_shared_qparams verifies that after prepare_pt2e the transpose node's input and output share the same observer instance.
1 parent d8d706a commit bc89f05

2 files changed

Lines changed: 67 additions & 0 deletions

File tree

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1045,6 +1045,11 @@ def _is_share_obs_or_fq_op(op: Callable) -> bool:
10451045
torch.ops.aten.mean.dim,
10461046
torch.ops.aten.permute.default,
10471047
torch.ops.aten.permute_copy.default,
1048+
torch.ops.aten.transpose.int,
1049+
torch.ops.aten.transpose_copy.int,
1050+
torch.ops.aten.t.default,
1051+
torch.ops.aten.t_copy.default,
1052+
torch.ops.aten.swapaxes.default,
10481053
torch.ops.aten.squeeze.dim,
10491054
torch.ops.aten.squeeze_copy.dim,
10501055
# TODO: remove?

backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -628,6 +628,68 @@ def forward(self, x):
628628
]
629629
self._test_quantizer(m, example_inputs, quantizer, node_occurrence, node_list)
630630

631+
def test_transpose_cat_shared_qparams(self):
632+
"""transpose feeding cat must share quantization parameters.
633+
634+
When a transpose output feeds into cat alongside a branch with a larger
635+
activation range, XNNPACK's cat annotator ties the transpose output to
636+
the combined-range scale while the transpose input keeps the upstream
637+
(narrow) scale. Without SharedQuantizationSpec on transpose the two
638+
scales diverge and XNNPACK rejects the model with
639+
xnn_status_invalid_parameter at load_method time.
640+
"""
641+
642+
class TransposeCat(torch.nn.Module):
643+
def __init__(self):
644+
super().__init__()
645+
self.l1 = torch.nn.Linear(16, 16)
646+
self.l2 = torch.nn.Linear(8, 8)
647+
648+
def forward(self, x, y):
649+
a = self.l1(x).transpose(1, 2) # (1, 16, 8) — narrow range
650+
b = self.l2(y) # (1, 16, 8) — large range
651+
return torch.cat([a, b], dim=1)
652+
653+
torch.manual_seed(0)
654+
quantizer = XNNPACKQuantizer()
655+
quantizer.set_global(get_symmetric_quantization_config())
656+
example_inputs = (torch.randn(1, 8, 16), torch.randn(1, 16, 8) * 8.0)
657+
m = TransposeCat().eval()
658+
m = export(m, example_inputs, strict=True).module()
659+
m = prepare_pt2e(m, quantizer)
660+
661+
torch.manual_seed(0)
662+
for _ in range(30):
663+
m(torch.randn(1, 8, 16), torch.randn(1, 16, 8) * 8.0)
664+
665+
m = convert_pt2e(m)
666+
667+
# After convert_pt2e, the Q/DQ flanking the transpose must carry
668+
# identical scale/zero_point; XNNStaticTranspose hard-requires this.
669+
for n in m.graph.nodes:
670+
if n.op != "call_function" or n.target != torch.ops.aten.transpose.int:
671+
continue
672+
dq_in = n.args[0]
673+
q_out = next(
674+
(
675+
u
676+
for u in n.users
677+
if u.op == "call_function" and "quantize" in str(u.target)
678+
),
679+
None,
680+
)
681+
self.assertIsNotNone(q_out, "expected a quantize node after transpose")
682+
self.assertEqual(
683+
dq_in.args[1],
684+
q_out.args[1],
685+
"transpose input/output scale must match for XNNPACK",
686+
)
687+
self.assertEqual(
688+
dq_in.args[2],
689+
q_out.args[2],
690+
"transpose input/output zero_point must match for XNNPACK",
691+
)
692+
631693
def test_propagate_annotation(self):
632694
quantizer = XNNPACKQuantizer()
633695
quantization_config = get_symmetric_quantization_config(is_per_channel=True)

0 commit comments

Comments
 (0)