[Relay] Support nchwc layout in ConvertLayout pass (#9681)

Author: Ever-Kid

python/tvm/relay/op/nn/_nn.py

@@ -17,6 +17,7 @@
 # pylint: disable=no-else-return, invalid-name, unused-argument, too-many-arguments, consider-using-in
 """Backend compiler related feature registration"""
 from __future__ import absolute_import
+import re
 
 from tvm import topi, relay
 from tvm.topi.utils import get_const_tuple
@@ -283,8 +284,9 @@ def convert_conv2d(attrs, inputs, tinfos, desired_layouts):
     desired_data_layout, desired_kernel_layout = map(str, desired_layouts)
     assert desired_data_layout != "default", "Data layout cannot be default"
     new_attrs["data_layout"] = desired_data_layout
+    need_tile = re.match(r"NCHW(\d*)c", desired_data_layout)
 
-    if desired_kernel_layout != "default":
+    if desired_kernel_layout != "default" and not need_tile:
         new_attrs["kernel_layout"] = desired_kernel_layout
         return relay.nn.conv2d(data, weight, **new_attrs)
 
@@ -309,6 +311,14 @@ def convert_conv2d(attrs, inputs, tinfos, desired_layouts):
     elif desired_data_layout == "HWNC":
         new_attrs["kernel_layout"] = "HWOI"
         return relay.nn.conv2d(data, weight, **new_attrs)
+    elif need_tile:
+        assert desired_kernel_layout != "default", "Kernel layout cannot be default."
+        tile = int(need_tile.group(1))
+        if isinstance(data, relay.expr.Var) and data.checked_type.shape[1] % tile != 0:
+            return relay.nn.conv2d(data, weight, **attrs)
+        else:
+            new_attrs["kernel_layout"] = desired_kernel_layout
+            return relay.nn.contrib_conv2d_nchwc(data, weight, **new_attrs)
 
     raise ValueError("Layout %s is not yet supported." % desired_data_layout)
 

tests/python/relay/test_pass_convert_op_layout.py

@@ -2323,5 +2323,116 @@ def expected():
     assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + "\n\n Expected = \n" + str(b)
 
 
+@pytest.mark.parametrize(
+    "data_layout, kernel_layout",
+    [
+        ("NCHW1c", "OIHW1i1o"),
+        ("NCHW4c", "OIHW4i4o"),
+        ("NCHW8c", "OIHW8i8o"),
+        ("NCHW16c", "OIHW16i16o"),
+    ],
+)
+def test_resnet_convert_layout_nchwc(data_layout, kernel_layout):
+    x = relay.var("x", shape=(1, 3, 224, 224))
+    weight1 = relay.var("weight1", shape=(64, 3, 7, 7))
+    weight2 = relay.var("weight2", shape=(64, 64, 3, 3))
+    weight3 = relay.var("weight3", shape=(64, 64, 1, 1))
+
+    def before():
+        y = relay.nn.conv2d(
+            x,
+            weight1,
+            strides=(2, 2),
+            padding=(3, 3),
+            channels=64,
+            kernel_size=(7, 7),
+            data_layout="NCHW",
+            kernel_layout="OIHW",
+        )
+        y = relay.nn.relu(y)
+        y = relay.nn.max_pool2d(y, pool_size=(3, 3), strides=(2, 2), padding=(1, 1))
+        y1 = relay.nn.conv2d(
+            y,
+            weight2,
+            channels=64,
+            kernel_size=(3, 3),
+            padding=(1, 1),
+            data_layout="NCHW",
+            kernel_layout="OIHW",
+        )
+        y1 = relay.nn.relu(y1)
+        y2 = relay.nn.conv2d(
+            y,
+            weight3,
+            channels=64,
+            kernel_size=(1, 1),
+            data_layout="NCHW",
+            kernel_layout="OIHW",
+        )
+        y2 = relay.nn.relu(y2)
+        y = y1 + y2
+        y = relay.nn.global_max_pool2d(y, layout="NCHW")
+        return y
+
+    def expected():
+        if data_layout == "NCHW1c":
+            y = relay.nn.contrib_conv2d_nchwc(
+                relay.layout_transform(x, "NCHW", data_layout),
+                relay.layout_transform(weight1, "OIHW", kernel_layout),
+                strides=(2, 2),
+                padding=(3, 3),
+                channels=64,
+                kernel_size=(7, 7),
+                data_layout=data_layout,
+                kernel_layout=kernel_layout,
+            )
+            y = relay.nn.relu(y)
+            y = relay.nn.max_pool2d(
+                y, pool_size=(3, 3), strides=(2, 2), padding=(1, 1), layout=data_layout
+            )
+        else:
+            y = relay.nn.conv2d(
+                x,
+                weight1,
+                strides=(2, 2),
+                padding=(3, 3),
+                channels=64,
+                kernel_size=(7, 7),
+                data_layout="NCHW",
+                kernel_layout="OIHW",
+            )
+            y = relay.nn.relu(y)
+            y = relay.nn.max_pool2d(y, pool_size=(3, 3), strides=(2, 2), padding=(1, 1))
+            y = relay.layout_transform(y, "NCHW", data_layout)
+        y1 = relay.nn.contrib_conv2d_nchwc(
+            y,
+            relay.layout_transform(weight2, "OIHW", kernel_layout),
+            channels=64,
+            kernel_size=(3, 3),
+            padding=(1, 1),
+            data_layout=data_layout,
+            kernel_layout=kernel_layout,
+        )
+        y1 = relay.nn.relu(y1)
+        y2 = relay.nn.contrib_conv2d_nchwc(
+            y,
+            relay.layout_transform(weight3, "OIHW", kernel_layout),
+            channels=64,
+            kernel_size=(1, 1),
+            data_layout=data_layout,
+            kernel_layout=kernel_layout,
+        )
+        y2 = relay.nn.relu(y2)
+        y = y1 + y2
+        y = relay.nn.global_max_pool2d(y, layout=data_layout)
+        y = relay.layout_transform(y, data_layout, "NCHW")
+        return y
+
+    a = before()
+    a = run_opt_pass(a, transform.ConvertLayout({"nn.conv2d": [data_layout, kernel_layout]}))
+    b = run_opt_pass(expected(), transform.InferType())
+    assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + "\n Expect = \n" + str(b)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])