Incompatible broadcast issue - Simple test case


#1
import tvm
from tvm import relay


p0 = relay.var('p0', shape=(1, 7, 7, 832), dtype='uint8')
p1 = relay.var('p1', shape=(1, 1, 832, 384), dtype='int8')
p2 = relay.var('p2', shape=(1, 7, 7, 384), dtype='int32')
conv = relay.nn.conv2d(p0, p1, kernel_size=(1, 1), data_layout='NHWC', kernel_layout='HWIO',
        out_dtype='int32')
subtract = relay.subtract(conv, p2)

z = subtract
func = relay.Function([p0, p1, p2], z)
mod = relay.Module.from_expr(func)
mod = relay.transform.InferType()(mod)
print(mod)

with relay.build_config(opt_level=3):
    graph, lib, params = relay.build(mod, 'llvm')

The above test fails with

TVMError: Check failed: false: Incompatible broadcast dims: 832 and 384 in: [1, 7, 7, 832] and [1, 7, 7, 384]

The print(mod) above shows well-formed Expr

v0.0.3
def @main(%p0: Tensor[(1, 7, 7, 832), uint8], %p1: Tensor[(1, 1, 832, 384), int8], %p2: Tensor[(1, 7, 7, 384), int32]) -> Tensor[(1, 7, 7, 384), int32] {
  %0 = nn.conv2d(%p0, %p1, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO", out_dtype="int32") /* ty=Tensor[(1, 7, 7, 384), int32] */;
  subtract(%0, %p2) /* ty=Tensor[(1, 7, 7, 384), int32] */
}

@zhiics
@vinx13
@yzhliu
@tqchen
@ziheng


#2
diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py
index 12b2c14..b629187 100644
--- a/topi/python/topi/x86/conv2d.py
+++ b/topi/python/topi/x86/conv2d.py
@@ -139,12 +139,12 @@ def _declaration_conv(cfg, data, kernel, strides, padding, dilation, layout, out
     kh, kw, _, _ = get_const_tuple(kernel.shape)
     if layout == 'HWCN':
         return nn.conv2d_hwcn(data, kernel, strides, padding, dilation, out_dtype)
-    elif layout == 'NHWC' and kh == 1 and kw == 1 and kernel.dtype == "int8":
-        if cfg.is_fallback:
-            _get_default_config(cfg, data, kernel, strides, padding, out_dtype, False, layout)
-        # specialize for INT8 1X1 conv on X86
-        return conv2d_avx_1x1._declaration_conv_nhwc_pack(cfg, data, kernel, strides,
-                                                          padding, dilation, out_dtype)
+    # elif layout == 'NHWC' and kh == 1 and kw == 1 and kernel.dtype == "int8":
+    #     if cfg.is_fallback:
+    #         _get_default_config(cfg, data, kernel, strides, padding, out_dtype, False, layout)
+    #     # specialize for INT8 1X1 conv on X86
+    #     return conv2d_avx_1x1._declaration_conv_nhwc_pack(cfg, data, kernel, strides,
+    #                                                       padding, dilation, out_dtype)
     elif layout == 'NHWC':
         return nn.conv2d_nhwc(data, kernel, strides, padding, dilation, out_dtype)
     raise ValueError("not support this layout {} yet".format(layout))

This is the workaround. Currently, the nhwc_pack assumes HWOI layout. We dont pass kernel_layout anywhere, causing this problem.