From b7208a85ab8db80ec1aaf8ccb7148b13a841a0ae Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sun, 2 Dec 2018 07:35:36 -0800 Subject: [PATCH] Fix topi conv (#29) --- topi/python/topi/arm_cpu/conv2d.py | 1 + topi/python/topi/x86/__init__.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/topi/python/topi/arm_cpu/conv2d.py b/topi/python/topi/arm_cpu/conv2d.py index 22c9d2368de3f..ebdf926ccbd75 100644 --- a/topi/python/topi/arm_cpu/conv2d.py +++ b/topi/python/topi/arm_cpu/conv2d.py @@ -210,6 +210,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 index 9e0e94e6cd2db..3641c0cd9c58e 100644 --- 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 *