Fix topi conv (apache#29)

tmoreau89 · Mar 22, 2019 · bcc6df6 · bcc6df6
1 parent 2c505d5
commit bcc6df6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/topi/python/topi/arm_cpu/conv2d.py b/topi/python/topi/arm_cpu/conv2d.py
@@ -211,6 +211,7 @@ def _decl_spatial_pack(cfg, data, kernel, strides, padding, dilation, layout, ou
     output = tvm.compute(oshape, lambda n, co, h, w:
                          conv[n][co//VC][h//VH][w//VW][h%VH][w%VW][co%VC],
                          name='output_unpack', tag='spatial_conv2d_output')
+    cfg.add_flop(2 * np.prod(oshape) * CI * KH * KW)
     return output
 
 def _schedule_spatial_pack(cfg, s, data_vec, kernel_vec,

diff --git a/topi/python/topi/x86/__init__.py b/topi/python/topi/x86/__init__.py
@@ -2,7 +2,10 @@
 """x86 specific declaration and schedules."""
 from __future__ import absolute_import as _abs
 
-from .conv2d import schedule_conv2d, schedule_conv2d_nhwc
+#from .conv2d import schedule_conv2d, schedule_conv2d_nhwc
+# TODO(lmzheng): don't know why the schedule call of arm_cpu will jump
+# to the compute call of x86 cpu. Some registerations are messed up.
+# Disable x86 registers temporarily for vta.
 from .binarize_pack import schedule_binarize_pack
 from .binary_dense import schedule_binary_dense
 from .nn import *