diff --git a/topi/python/topi/cuda/conv2d_winograd.py b/topi/python/topi/cuda/conv2d_winograd.py index 50c485290bbf..eb004d79fd9b 100644 --- a/topi/python/topi/cuda/conv2d_winograd.py +++ b/topi/python/topi/cuda/conv2d_winograd.py @@ -17,6 +17,7 @@ # pylint: disable=invalid-name,unused-variable,unused-argument """Winograd template for cuda backend""" +import logging import tvm from tvm import autotvm @@ -27,6 +28,8 @@ from ..nn.winograd_util import winograd_transform_matrices +logger = logging.getLogger('conv2d_winograd') + def _infer_tile_size(data, kernel): N, CI, H, W = get_const_tuple(data.shape) @@ -42,26 +45,25 @@ def winograd_cuda(cfg, data, kernel, strides, padding, dilation, layout, out_dty N, CI, H, W = get_const_tuple(data.shape) + if isinstance(dilation, int): + dilation_h = dilation_w = dilation + else: + dilation_h, dilation_w = dilation + HSTR, WSTR = (strides, strides) if isinstance(strides, int) else strides + if not pre_computed: # kernel tensor is raw tensor, do strict check - if isinstance(dilation, int): - dilation_h = dilation_w = dilation - else: - dilation_h, dilation_w = dilation if dilation_h != 1 or dilation_w != 1: - kernel = dilate(kernel, (1, 1, dilation_h, dilation_w)) - + kernel = dilation(kernel, (1, 1, dilation_h, dilation_w)) CO, CI, KH, KW = get_const_tuple(kernel.shape) - HPAD, WPAD, _, _ = nn.get_pad_tuple(padding, kernel) - HSTR, WSTR = (strides, strides) if isinstance(strides, int) else strides assert HSTR == 1 and WSTR == 1 and KH == KW - else: # kernel tensor is pre-transfomred. this op is created by - # alter op layout, do not check + else: + # kernel tensor is pre-transfomred. this op is created by alter op layout. # dilation is not supported - HSTR = WSTR = 1 - HPAD = WPAD = 1 - KH = KW = 3 _, _, CI, CO = get_const_tuple(kernel.shape) + KH = KW = 3 + assert HSTR == 1 and WSTR == 1 and dilation_h == 1 and dilation_w == 1 + HPAD, WPAD, _, _ = nn.get_pad_tuple(padding, kernel) data_pad = nn.pad(data, (0, 0, HPAD, WPAD), (0, 0, HPAD, WPAD), name="data_pad") r = KW @@ -384,7 +386,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, F): return F.nn.conv2d(*copy_inputs, **new_attrs) if attrs.get_int_tuple("dilation") != (1, 1): - warnings.warn("Does not support weight pre-transform for dilated convolution.") + logger.warning("Does not support weight pre-transform for dilated convolution.") return None # pre-compute weight transformation in winograd diff --git a/topi/tests/python/common.py b/topi/tests/python/common.py index 0bacd61129b7..4e0a45be0a22 100644 --- a/topi/tests/python/common.py +++ b/topi/tests/python/common.py @@ -40,4 +40,5 @@ def _query_inside(self, target, workload): cfg = FallbackConfigEntity() cfg.template_key = 'int8' self.memory[key] = cfg + cfg.is_fallback = False return cfg diff --git a/topi/tests/python/test_topi_conv2d_winograd.py b/topi/tests/python/test_topi_conv2d_winograd.py index 908317ebf55e..5974dad20f88 100644 --- a/topi/tests/python/test_topi_conv2d_winograd.py +++ b/topi/tests/python/test_topi_conv2d_winograd.py @@ -99,6 +99,7 @@ def _query_inside(self, target, workload): cfg = FallbackConfigEntity() cfg.template_key = 'winograd' self.memory[key] = cfg + cfg.is_fallback = False return cfg