Fuse view/tranpose/view as shuffle

Author: EikanWang

tests/cpu/test_jit.py

@@ -235,6 +235,22 @@ def forward(self, x):
         c = torch.add(b, b)
         return c
 
+class ChannelShuffle(nn.Module):
+    def __init__(self, batchsize, num_channels, height, width, groups):
+        super(ChannelShuffle, self).__init__()
+        self.batchsize = batchsize
+        self.num_channels = num_channels
+        self.height = height
+        self.width = width
+        self.groups = groups
+
+    def forward(self, x):
+        channels_per_group = self.num_channels // self.groups
+        x = x.view(self.batchsize, self.groups, channels_per_group, self.height, self.width)
+        x = torch.transpose(x, 1, 2).contiguous()
+        x = x.view(self.batchsize, -1, self.height, self.width)
+        return x
+
 
 class Tester(TestCase):
 
@@ -528,6 +544,13 @@ def test_output_linear_relu(self):
             kind_in_graph="ipex::linear_relu")
 
 
+    def test_channel_shuffle(self):
+        self._test_output(
+            ChannelShuffle(10, 16, 50, 50, 4),
+            torch.rand(10, 16, 50, 50),
+            kind_in_graph="ipex::shuffle_2d")
+
+
     def test_jit_function(self):
         # test hool trace and script can works for function
         def fn(input, weight, bias):

torch_ipex/csrc/jit/graph_rewrite.cpp

@@ -55,11 +55,11 @@ std::unordered_map<std::string, c10::IValue> getConvParams(
 void FuseShuffle(std::shared_ptr<Graph>& graph) {
   std::string shuffle = R"(
       graph(%input, %view_shape:int[], %trans_dim0:int, %trans_dim1:int, %mem_format:int, %flattern_shape:int[]):
-        %r = aten::view(%input, %view_shape)
-        %r = aten::transpose(%r, %trans_dim0, %trans_dim1)
-        %r = aten::contiguous(%r, %mem_format)
-        %r = aten::view(%r, %flattern_shape)
-        return (%r) )";
+        %r1 = aten::view(%input, %view_shape)
+        %r2 = aten::transpose(%r1, %trans_dim0, %trans_dim1)
+        %r3 = aten::contiguous(%r2, %mem_format)
+        %r4 = aten::view(%r3, %flattern_shape)
+        return (%r4) )";
 
   std::string shuffle_2d_fusion = R"(
       graph(%input, %view_shape:int[], %trans_dim0:int, %trans_dim1:int, %mem_format:int, %flattern_shape:int[]):