Skip to content

Commit 8384872

Browse files
Arm backend: fix Torch compatibility
- Recent pin update to Torch 2.12 no longer exposes .out overloads for quantized-decomposed operator. - Code for backward-compatiliby is included, but should be removed in the future when older Torch versions are no longer supported. Signed-off-by: Sangwon Ha <sangwon.ha@arm.com> Change-Id: I9b8c76cbfdba36b0b17a772a66a785fdc920f347
1 parent 769b0cf commit 8384872

1 file changed

Lines changed: 35 additions & 6 deletions

File tree

backends/arm/test/runner_utils.py

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,40 @@
4545

4646
logger = logging.getLogger(__name__)
4747

48+
# For Torch 2.12 and later overloads.
49+
_QDQ_TORCH_OVERLOADS = (
50+
("quantize_per_tensor", ("tensor", "tensor2", "default")),
51+
("dequantize_per_tensor", ("tensor", "tensor2", "default")),
52+
("quantize_per_channel", ("default",)),
53+
("dequantize_per_channel", ("default",)),
54+
)
55+
56+
# For backward compatibility with Torch versions older than 2.12.
57+
_QDQ_BACKWARD_COMPAT_OVERLOADS = (
58+
("quantize_per_tensor", ("out",)),
59+
("dequantize_per_tensor", ("out",)),
60+
("quantize_per_channel", ("out",)),
61+
("dequantize_per_channel", ("out",)),
62+
)
63+
64+
65+
def _get_qdq_memory_format_ops() -> tuple[object, ...]:
66+
qdq_ops = []
67+
backward_compat = dict(_QDQ_BACKWARD_COMPAT_OVERLOADS)
68+
ns = torch.ops.quantized_decomposed
69+
for op_name, overload_names in _QDQ_TORCH_OVERLOADS:
70+
op_packet = getattr(ns, op_name, None)
71+
if op_packet is None:
72+
continue
73+
for overload_name in overload_names + backward_compat[op_name]:
74+
if hasattr(op_packet, overload_name):
75+
qdq_ops.append(getattr(op_packet, overload_name))
76+
77+
return tuple(qdq_ops)
78+
79+
80+
_QDQ_MEMORY_FORMAT_OPS = _get_qdq_memory_format_ops()
81+
4882
# Copied from PyTorch.
4983
# From torch/testing/_internal/common_utils.py:torch_to_numpy_dtype_dict
5084
# To avoid a dependency on _internal stuff.
@@ -358,12 +392,7 @@ def __torch_function__(self, func, types, args=..., kwargs=None):
358392

359393
# This is a hack since Q/DQ ops does not handle channels last input correctly: the simplest and most robust
360394
# workaround is to simply run them in channels first format and then convert back to channels last.
361-
if func in (
362-
torch.ops.quantized_decomposed.quantize_per_tensor.out,
363-
torch.ops.quantized_decomposed.dequantize_per_tensor.out,
364-
torch.ops.quantized_decomposed.quantize_per_channel.out,
365-
torch.ops.quantized_decomposed.dequantize_per_channel.out,
366-
):
395+
if func in _QDQ_MEMORY_FORMAT_OPS:
367396

368397
input_dim_order = args[0].dim_order()
369398
if input_dim_order in (NHWC_ORDER, NNHWC_ORDER):

0 commit comments

Comments
 (0)