From c151f9095ee5031117257ac0bd39f83a01992141 Mon Sep 17 00:00:00 2001 From: Giuseppe Rossini Date: Mon, 15 Jun 2020 11:29:59 +0100 Subject: [PATCH] Address review comments Change-Id: Ia1755a0af7b6d159072d9f0c93c932c481101e48 --- python/tvm/relay/op/nn/_nn.py | 1 - topi/python/topi/arm_cpu/conv2d_alter_op.py | 5 +++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index 53f317dcccfee..a4381cb9ae60e 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -452,7 +452,6 @@ def compute_mirror_pad(attrs, inputs, out_dtype): reg.register_pattern("nn.contrib_conv2d_gemm_without_weight_transform", OpPattern.OUT_ELEMWISE_FUSABLE) - @reg.register_compute("nn.contrib_conv2d_gemm_weight_transform") def compute_contrib_conv2d_gemm_weight_transform(attrs, inputs, out_dtype): """Compute definition of contrib_conv2d_gemm_weight_transform""" diff --git a/topi/python/topi/arm_cpu/conv2d_alter_op.py b/topi/python/topi/arm_cpu/conv2d_alter_op.py index c23879a0c664a..085a84d819adb 100644 --- a/topi/python/topi/arm_cpu/conv2d_alter_op.py +++ b/topi/python/topi/arm_cpu/conv2d_alter_op.py @@ -91,6 +91,8 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type): assert (data.dtype == 'int8' and kernel.dtype == 'int8' or data.dtype == 'uint8' and kernel.dtype == 'uint8') + assert data_layout == "NHWC" and kernel_layout == "HWIO" + data_expr, kernel_expr = inputs data_int16 = relay.cast(data_expr, dtype='int16') @@ -242,6 +244,9 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type): if topi_tmpl == "conv2d_NHWC_quantized.arm_cpu": assert (data.dtype == 'int8' and kernel.dtype == 'int8' or data.dtype == 'uint8' and kernel.dtype == 'uint8') + + assert data_layout == "NHWC" and kernel_layout == "HWIO" + CO, IC, KH, KW = get_const_tuple(kernel.shape) K = KH * KW * IC