[quant][graphmode][refactor] Support filter function in quant fusion patterns (#35333)

jerryzh168 · facebook-github-bot · commit e90c32f11fc5 · 2020-03-28T08:23:44.000-07:00
Summary: Pull Request resolved: #35333 Test Plan: regression tests in: python test/test_jit.py Imported from OSS Differential Revision: D20655312 fbshipit-source-id: 50b937bc56aff93f20fe9a0079bf3aec50f6d25d
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -2473,10 +2473,10 @@ void SwapDeQuant(std::shared_ptr<Graph>& graph) {
 }
 
 void QuantFusion(std::shared_ptr<Graph>& graph) {
-  for (const auto& item : quant_fusion_pattern_and_replacements()) {
+  for (const auto& info : quant_fusion_pattern_and_replacements()) {
     SubgraphRewriter rewriter;
-    rewriter.RegisterRewritePattern(item.first, item.second);
-    rewriter.runOnGraph(graph);
+    rewriter.RegisterRewritePattern(info.pattern, info.replacement);
+    rewriter.runOnGraph(graph, info.filter);
   }
 }
 
diff --git a/torch/csrc/jit/passes/quantization_patterns.h b/torch/csrc/jit/passes/quantization_patterns.h
@@ -2,12 +2,23 @@
 
 #include <string>
 #include <unordered_map>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/subgraph_matcher.h>
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
 
 namespace torch {
 namespace jit {
 
-std::unordered_map<std::string, std::string>
-quant_fusion_pattern_and_replacements() {
+struct QuantFusionInfo {
+  std::string quantized_op_name;
+  std::string pattern;
+  std::string replacement;
+  std::function<bool(const Match&, const std::unordered_map<std::string, Value*>&)> filter = [](const Match&, const std::unordered_map<std::string, Value*>&) {
+     return true;
+  };
+};
+
+std::vector<QuantFusionInfo> quant_fusion_pattern_and_replacements() {
   std::string conv2d = R"(
 graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
         %a_dequant = aten::dequantize(%a_quant)
@@ -165,18 +176,18 @@ graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
   // We don't have quantized inplace add right now
 
   return {
-      {conv2d, quantized_conv2d},
-      {conv2d_relu, quantized_conv2d_relu},
-      {conv2d_inplace_relu, quantized_conv2d_relu},
-      {addmm, quantized_linear},
-      {matmul_with_bias, quantized_linear},
-      {matmul_no_bias, quantized_linear_no_bias},
-      {aten_linear, quantized_aten_linear},
-      {add_relu, quantized_add_relu},
-      {add_inplace_relu, quantized_add_relu},
-      {add, quantized_add},
-      {inplace_add, quantized_add},
-      {cat, quantized_cat},
+      {"quantized::conv2d", conv2d, quantized_conv2d},
+      {"quantized::conv2d_relu", conv2d_relu, quantized_conv2d_relu},
+      {"quantized::conv2d_relu", conv2d_inplace_relu, quantized_conv2d_relu},
+      {"quantized::linear", addmm, quantized_linear},
+      {"quantized::linear", matmul_with_bias, quantized_linear},
+      {"quantized::linear", matmul_no_bias, quantized_linear_no_bias},
+      {"quantized::linear", aten_linear, quantized_aten_linear},
+      {"quantized::add_relu", add_relu, quantized_add_relu},
+      {"quantized::add_relu", add_inplace_relu, quantized_add_relu},
+      {"quantized::add", add, quantized_add},
+      {"quantized::add", inplace_add, quantized_add},
+      {"quantized::cat", cat, quantized_cat},
   };
 }
 

Original file line number	Diff line number	Diff line change
`@@ -2473,10 +2473,10 @@ void SwapDeQuant(std::shared_ptr<Graph>& graph) {`
`2473`	`2473`	`}`
`2474`	`2474`
`2475`	`2475`	`void QuantFusion(std::shared_ptr<Graph>& graph) {`
`2476`		`- for (const auto& item : quant_fusion_pattern_and_replacements()) {`
	`2476`	`+ for (const auto& info : quant_fusion_pattern_and_replacements()) {`
`2477`	`2477`	`SubgraphRewriter rewriter;`
`2478`		`- rewriter.RegisterRewritePattern(item.first, item.second);`
`2479`		`- rewriter.runOnGraph(graph);`
	`2478`	`+ rewriter.RegisterRewritePattern(info.pattern, info.replacement);`
	`2479`	`+ rewriter.runOnGraph(graph, info.filter);`
`2480`	`2480`	`}`
`2481`	`2481`	`}`
`2482`	`2482`