fix group conv3d pack kernel shape error (#12523)

apache · Aug 22, 2022 · e9aad35 · e9aad35
1 parent 2629065
commit e9aad35
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/python/tvm/topi/x86/conv3d.py b/python/tvm/topi/x86/conv3d.py
@@ -272,10 +272,14 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dty
         shape, lambda n, C, d, h, c, w: data_pad[n, d, h, w, C * ic_bn + c], name="data_vec"
     )
 
+    ci_tile = in_channel // groups // ic_bn
+    if ci_tile == 0 or ci_tile * ic_bn * groups < in_channel:
+        ci_tile += 1
+
     # pack kernel
     shape = (
         num_filter // oc_bn,
-        in_channel // groups // ic_bn if (in_channel // groups // ic_bn) else 1,
+        ci_tile,
         kernel_depth,
         kernel_height,
         kernel_width,
@@ -389,10 +393,14 @@ def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, groups,
         shape, lambda n, C, d, h, c, w: data_pad[n, C * ic_bn + c, d, h, w], name="data_vec"
     )
 
+    ci_tile = in_channel // groups // ic_bn
+    if ci_tile == 0 or ci_tile * ic_bn * groups < in_channel:
+        ci_tile += 1
+
     # pack kernel
     shape = (
         num_filter // oc_bn,
-        in_channel // groups // ic_bn if (in_channel // groups // ic_bn) else 1,
+        ci_tile,
         kernel_depth,
         kernel_height,
         kernel_width,

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
@@ -2980,6 +2980,17 @@ def repeat(num, dims):
             group=8,
         )
 
+        verify_conv(
+            (1, 12) + repeat(5, dims),
+            (30, 4) + repeat(3, dims),
+            (1, 30) + repeat(5, dims),
+            2 * repeat(1, dims),
+            repeat(3, dims),
+            repeat(1, dims),
+            repeat(1, dims),
+            group=3,
+        )
+
 
 @tvm.testing.parametrize_targets
 def test_convtranspose(target, dev):