[pt2e][xnnpack_quantizer] quantize "mul" (pytorch#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: pytorch#110428 Approved by: https://github.com/jerryzh168 ghstack dependencies: pytorch#110427
yuguo68 · Oct 4, 2023 · 08c7dcd · 08c7dcd
1 parent 66202ed
commit 08c7dcd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
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 = []