fix conv_folding cannot capture Conv -> Mul -> Add etc (#602)

* fix conv_folding cannot capture Conv -> Mul -> Add etc

Author: blzheng

intel_extension_for_pytorch/csrc/jit/cpu/passes/frozen_conv_folding.cpp

@@ -119,10 +119,11 @@ bool checkConvAndBroadcastingOpPreConditions(Node* conv, Node* op) {
   return true;
 }
 
-void FoldFrozenConvAddOrSub(Block* b) {
+bool FoldFrozenConvAddOrSub(Block* b) {
+  bool graph_modified = false;
   for (Node* n : b->nodes()) {
     for (Block* block : n->blocks()) {
-      FoldFrozenConvAddOrSub(block);
+      graph_modified |= FoldFrozenConvAddOrSub(block);
     }
 
     if (supportedAddOrSub(n) && supportedConvNode(n->inputs().at(0)->node())) {
@@ -174,15 +175,18 @@ void FoldFrozenConvAddOrSub(Block* b) {
           add_or_sub->kind().toUnqualString());
       conv->replaceInputWith(conv_b_value, fused_conv_b);
       add_or_sub->output()->replaceAllUsesWith(conv->output());
+      graph_modified = true;
       // DCE run after cleans up nodes
     }
   }
+  return graph_modified;
 }
 
-void FoldFrozenConvMulOrDiv(Block* b) {
+bool FoldFrozenConvMulOrDiv(Block* b) {
+  bool graph_modified = false;
   for (Node* n : b->nodes()) {
     for (Block* block : n->blocks()) {
-      FoldFrozenConvMulOrDiv(block);
+      graph_modified |= FoldFrozenConvMulOrDiv(block);
     }
 
     if (supportedMulOrDiv(n) && supportedConvNode(n->inputs().at(0)->node())) {
@@ -287,21 +291,35 @@ void FoldFrozenConvMulOrDiv(Block* b) {
             mul_or_div->kind().toUnqualString());
         conv->replaceInputWith(conv_b_value, fused_conv_bias);
       }
+      graph_modified = true;
       // DCE run after cleans up nodes
     }
   }
+  return graph_modified;
 }
 
 } // namespace
 
-void FoldFrozenConvAddOrSub(std::shared_ptr<Graph>& graph) {
-  FoldFrozenConvAddOrSub(graph->block());
+bool FoldFrozenConvAddOrSub(std::shared_ptr<Graph>& graph) {
+  bool graph_modified = FoldFrozenConvAddOrSub(graph->block());
   EliminateDeadCode(graph);
+  return graph_modified;
 }
 
-void FoldFrozenConvMulOrDiv(std::shared_ptr<Graph>& graph) {
-  FoldFrozenConvMulOrDiv(graph->block());
+bool FoldFrozenConvMulOrDiv(std::shared_ptr<Graph>& graph) {
+  bool graph_modified = FoldFrozenConvMulOrDiv(graph->block());
   EliminateDeadCode(graph);
+  return graph_modified;
+}
+
+void FrozenConvFolding(std::shared_ptr<Graph>& graph) {
+  // run a couple times to capture Conv -> Mul -> Add etc
+  bool changed;
+  do {
+    changed = false;
+    changed |= FoldFrozenConvAddOrSub(graph);
+    changed |= FoldFrozenConvMulOrDiv(graph);
+  } while (changed);
 }
 
 } // namespace jit

intel_extension_for_pytorch/csrc/jit/cpu/passes/frozen_conv_folding.h

@@ -8,12 +8,15 @@ namespace jit {
 // Fuses Convolution -> Add/Sub into a single Convolution by
 // folding add constant tensor into conv weights.
 // This pass only works on Frozen Graphs; otherwise it is a No-Op.
-TORCH_API void FoldFrozenConvAddOrSub(std::shared_ptr<Graph>& graph);
+TORCH_API bool FoldFrozenConvAddOrSub(std::shared_ptr<Graph>& graph);
 
 // Fuses Convolution -> Mul/Div into a single Convolution by
 // folding add constant tensor into conv weights.
 // This pass only works on Frozen Graphs; otherwise it is a No-Op.
-TORCH_API void FoldFrozenConvMulOrDiv(std::shared_ptr<Graph>& graph);
+TORCH_API bool FoldFrozenConvMulOrDiv(std::shared_ptr<Graph>& graph);
+
+// Call FoldFrozenConvAddOrSub and FoldFrozenConvMulOrDiv a couple times
+TORCH_API void FrozenConvFolding(std::shared_ptr<Graph>& graph);
 
 } // namespace jit
 } // namespace torch

intel_extension_for_pytorch/csrc/jit/fusion_pass.cpp

@@ -359,8 +359,7 @@ void IPEXFusionPass(std::shared_ptr<Graph>& graph) {
   graph_rewrite_helper::replaceConvolutionWithAtenConv(graph);
 
   // convolution folding
-  FoldFrozenConvAddOrSub(graph);
-  FoldFrozenConvMulOrDiv(graph);
+  FrozenConvFolding(graph);
 
   // convolution fusion
   graph_rewrite::insertPrePackedConvOp(graph);

tests/cpu/test_jit.py

@@ -74,6 +74,13 @@
 from torch._six import inf, nan
 from torch.testing._internal.common_utils import TestCase
 
+try:
+    import torchvision
+    HAS_TORCHVISION = True
+except ImportError:
+    HAS_TORCHVISION = False
+skipIfNoTorchVision = unittest.skipIf(not HAS_TORCHVISION, "no torchvision")
+
 device = 'cpu:0'
 SIZE = 100
 
@@ -2508,6 +2515,55 @@ def test_remove_bailout(self):
                 kind_not_in_graph="prim::BailOut",
                 prec=0.05)
 
+    @skipIfNoTorchVision
+    def test_conv_torchvision_bn_folding(self):
+        from torchvision.ops import misc as misc_nn_ops
+        class M(nn.Module):
+            def __init__(self):
+                super(M, self).__init__()
+                norm_layer = misc_nn_ops.FrozenBatchNorm2d
+                self.inplanes = 64
+                self.dilation = 1
+                self.groups = 1
+                self.base_width = 64
+                self.conv1 = torch.nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+                self.bn1 = norm_layer(self.inplanes)
+                self.relu = torch.nn.ReLU(inplace=True)
+                self.maxpool = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+            def forward(self, x):
+                x = self.conv1(x)
+                x = self.bn1(x)
+                x = self.relu(x)
+                x = self.maxpool(x)
+                return x
+
+        model = M().eval()
+        self._test_output(
+            model,
+            torch.randn(1, 3, 1200, 1200),
+            kind_in_graph="ipex_prepack::convolution_relu_run",
+            kind_not_in_graph="aten::add")
+
+        self._test_output(
+            model,
+            torch.randn(1, 3, 1200, 1200),
+            kind_in_graph="ipex_prepack::convolution_relu_run",
+            kind_not_in_graph="aten::mul")
+
+        self._test_output_bf16(
+            model,
+            torch.randn(1, 3, 1200, 1200),
+            kind_in_graph="ipex_prepack::convolution_relu_run",
+            kind_not_in_graph="aten::add",
+            prec=0.1)
+
+        self._test_output_bf16(
+            model,
+            torch.randn(1, 3, 1200, 1200),
+            kind_in_graph="ipex_prepack::convolution_relu_run",
+            kind_not_in_graph="aten::mul",
+            prec=0.1)
 
 if __name__ == '__main__':
     torch.manual_seed(2020)