[pt2e][xnnpack_quantizer] quantize "mul" (#110428)

mcr229 · pytorchmergebot · commit 08c7dcda65ce · 2023-10-04T05:11:53.000Z
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: #110428 Approved by: https://github.com/jerryzh168 ghstack dependencies: #110427
diff --git 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
@@ -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 = []
