From 08c7dcda65cefd21df771f20ac55cc3846663a9c Mon Sep 17 00:00:00 2001 From: Max Ren Date: Tue, 3 Oct 2023 15:06:00 -0700 Subject: [PATCH] [pt2e][xnnpack_quantizer] quantize "mul" (#110428) Adding "mul" to list of partitions that are supported by the quantizer. This shows up in EDSR, where we still want to quantize the mul op Differential Revision: [D49850151](https://our.internmc.facebook.com/intern/diff/D49850151/) Pull Request resolved: https://github.com/pytorch/pytorch/pull/110428 Approved by: https://github.com/jerryzh168 ghstack dependencies: #110427 --- torch/ao/quantization/pt2e/graph_utils.py | 2 +- torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/torch/ao/quantization/pt2e/graph_utils.py b/torch/ao/quantization/pt2e/graph_utils.py index 2d4e73dee3577c..2390bafed28681 100644 --- a/torch/ao/quantization/pt2e/graph_utils.py +++ b/torch/ao/quantization/pt2e/graph_utils.py @@ -25,7 +25,7 @@ {torch.nn.BatchNorm2d, torch.nn.functional.batch_norm}, {torch.nn.Hardtanh, torch.nn.functional.hardtanh, torch.nn.functional.hardtanh_}, {torch.add, operator.add, operator.iadd, "add", "add_"}, - {torch.mul, operator.mul, operator.imul}, + {torch.mul, operator.mul, operator.imul, "mul", "mul_"}, ] diff --git a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py index 085c65b768bce0..e89c06715b2e16 100644 --- a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py +++ b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py @@ -739,7 +739,7 @@ def _annotate_mul( filter_fn: Optional[Callable[[Node], bool]] = None, ) -> Optional[List[List[Node]]]: mul_partitions = get_source_partitions( - gm.graph, [operator.mul, torch.mul, operator.imul], filter_fn + gm.graph, ["mul", "mul_", operator.mul, torch.mul, operator.imul], filter_fn ) mul_partitions = list(itertools.chain(*mul_partitions.values())) annotated_partitions = []