|
45 | 45 |
|
46 | 46 | logger = logging.getLogger(__name__) |
47 | 47 |
|
| 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 | + |
48 | 82 | # Copied from PyTorch. |
49 | 83 | # From torch/testing/_internal/common_utils.py:torch_to_numpy_dtype_dict |
50 | 84 | # To avoid a dependency on _internal stuff. |
@@ -358,12 +392,7 @@ def __torch_function__(self, func, types, args=..., kwargs=None): |
358 | 392 |
|
359 | 393 | # This is a hack since Q/DQ ops does not handle channels last input correctly: the simplest and most robust |
360 | 394 | # 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: |
367 | 396 |
|
368 | 397 | input_dim_order = args[0].dim_order() |
369 | 398 | if input_dim_order in (NHWC_ORDER, NNHWC_ORDER): |
|
0 commit comments