Address review comments

Change-Id: Ia1755a0af7b6d159072d9f0c93c932c481101e48
apache · Jun 16, 2020 · c151f90 · c151f90
1 parent e323d88
commit c151f90
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git 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
@@ -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