From c4d157984721570e6b32c7d2ae3f2c7ad53646c9 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Thu, 30 May 2019 15:01:52 -0700 Subject: [PATCH 1/8] Fix x86 depthwise conv2d alter_op_layout --- topi/python/topi/arm_cpu/depthwise_conv2d.py | 6 +++--- topi/python/topi/x86/conv2d.py | 3 ++- topi/python/topi/x86/depthwise_conv2d.py | 11 +++++++++-- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/topi/python/topi/arm_cpu/depthwise_conv2d.py b/topi/python/topi/arm_cpu/depthwise_conv2d.py index e09e355ad8aa..e500239d9d2c 100644 --- a/topi/python/topi/arm_cpu/depthwise_conv2d.py +++ b/topi/python/topi/arm_cpu/depthwise_conv2d.py @@ -26,11 +26,11 @@ from ..nn.util import get_pad_tuple # register original implementation of depthwise_conv2d_nchw since we don't need to change this part -autotvm.register_topi_compute(depthwise_conv2d_nchw, ['arm_cpu', 'cpu'], 'direct', +autotvm.register_topi_compute(depthwise_conv2d_nchw, ['arm_cpu'], 'direct', depthwise_conv2d_nchw.fdefault) # register customized schedule for arm cpu. -@autotvm.register_topi_schedule(schedule_depthwise_conv2d_nchw, ['arm_cpu', 'cpu'], +@autotvm.register_topi_schedule(schedule_depthwise_conv2d_nchw, ['arm_cpu'], ['direct', 'contrib_spatial_pack']) def schedule_depthwise_conv2d_nchw_arm(cfg, outs): """Schedule depthwise conv2d @@ -151,7 +151,7 @@ def _callback(op): traverse_inline(s, outs[0].op, _callback) return s -@autotvm.register_topi_compute(depthwise_conv2d_nchw, ['arm_cpu', 'cpu'], ['contrib_spatial_pack']) +@autotvm.register_topi_compute(depthwise_conv2d_nchw, ['arm_cpu'], ['contrib_spatial_pack']) def depthwise_conv2d_arm_cpu(cfg, data, kernel, strides, padding, dilation, out_dtype): """TOPI compute callback for depthwise_conv2d nchw diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py index d0894ad5a1b2..8a709b8af1e3 100644 --- a/topi/python/topi/x86/conv2d.py +++ b/topi/python/topi/x86/conv2d.py @@ -415,7 +415,8 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F): dtype = data.dtype out_dtype = dtype if out_dtype in ("same", "") else out_dtype - is_depthwise = groups == in_channel and groups == out_channel + kshape = get_const_tuple(kernel.shape) + is_depthwise = groups == kshape[0] and kshape[1] == 1 # only optimize for NCHW if layout != 'NCHW': diff --git a/topi/python/topi/x86/depthwise_conv2d.py b/topi/python/topi/x86/depthwise_conv2d.py index 6ea11f234759..ddcd8415df77 100644 --- a/topi/python/topi/x86/depthwise_conv2d.py +++ b/topi/python/topi/x86/depthwise_conv2d.py @@ -22,11 +22,12 @@ from tvm.autotvm.task.space import SplitEntity from tvm.autotvm.task.topi_integration import deserialize_args from .. import generic, tag +from ..generic import schedule_depthwise_conv2d_nchw from ..nn.pad import pad from ..util import get_const_tuple from ..nn.util import get_pad_tuple -from ..nn.depthwise_conv2d import depthwise_conv2d_NCHWc, _get_workload, \ - depthwise_conv2d_infer_layout +from ..nn.depthwise_conv2d import depthwise_conv2d_nchw, depthwise_conv2d_NCHWc, \ + _get_workload, depthwise_conv2d_infer_layout from .util import get_fp32_len @@ -70,6 +71,12 @@ def _fallback_schedule(cfg, wkl): cfg["tile_ow"] = SplitEntity([out_width // reg_n, reg_n]) +autotvm.register_topi_compute(depthwise_conv2d_nchw, 'cpu', 'direct', + depthwise_conv2d_nchw.fdefault) +autotvm.register_topi_schedule(schedule_depthwise_conv2d_nchw, 'cpu', 'direct', + schedule_depthwise_conv2d_nchw.fdefault) + + @autotvm.register_topi_compute(depthwise_conv2d_NCHWc, 'cpu', 'direct') def _depthwise_conv2d_NCHWc_cpu(cfg, data, kernel, strides, padding, dilation, layout, out_layout, out_dtype=None): From b285ee6afee22103625bdcde72b6e7d0d27273c5 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Thu, 30 May 2019 17:17:34 -0700 Subject: [PATCH 2/8] Small fix --- topi/python/topi/arm_cpu/depthwise_conv2d.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/topi/python/topi/arm_cpu/depthwise_conv2d.py b/topi/python/topi/arm_cpu/depthwise_conv2d.py index e500239d9d2c..51088df905ca 100644 --- a/topi/python/topi/arm_cpu/depthwise_conv2d.py +++ b/topi/python/topi/arm_cpu/depthwise_conv2d.py @@ -26,11 +26,11 @@ from ..nn.util import get_pad_tuple # register original implementation of depthwise_conv2d_nchw since we don't need to change this part -autotvm.register_topi_compute(depthwise_conv2d_nchw, ['arm_cpu'], 'direct', +autotvm.register_topi_compute(depthwise_conv2d_nchw, 'arm_cpu', 'direct', depthwise_conv2d_nchw.fdefault) # register customized schedule for arm cpu. -@autotvm.register_topi_schedule(schedule_depthwise_conv2d_nchw, ['arm_cpu'], +@autotvm.register_topi_schedule(schedule_depthwise_conv2d_nchw, 'arm_cpu', ['direct', 'contrib_spatial_pack']) def schedule_depthwise_conv2d_nchw_arm(cfg, outs): """Schedule depthwise conv2d @@ -151,7 +151,7 @@ def _callback(op): traverse_inline(s, outs[0].op, _callback) return s -@autotvm.register_topi_compute(depthwise_conv2d_nchw, ['arm_cpu'], ['contrib_spatial_pack']) +@autotvm.register_topi_compute(depthwise_conv2d_nchw, 'arm_cpu', ['contrib_spatial_pack']) def depthwise_conv2d_arm_cpu(cfg, data, kernel, strides, padding, dilation, out_dtype): """TOPI compute callback for depthwise_conv2d nchw From 57fe10e54e1e77092ccf71f90b4b6ad84442f110 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Fri, 31 May 2019 12:54:46 -0700 Subject: [PATCH 3/8] Add test case --- .../python/relay/test_pass_alter_op_layout.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index 2eea1c4ca87a..1a0f29cb6e80 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -559,6 +559,46 @@ def expected(): assert(alpha_equal(a, b)) +def test_alter_layout_depthwise_conv2d(): + """Test depthwise_conv2d operator""" + def before(): + x = relay.var("x", shape=(1, 32, 56, 56)) + w = relay.var("w", shape=(32, 1, 3, 3)) + y = relay.nn.conv2d(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), groups=32) + y = relay.Function(free_vars(y), y) + return y + + import topi + @register_alter_op_layout("nn.conv2d", level=111) + def alter_conv2d(attrs, inputs, tinfos): + return topi.x86.conv2d._alter_conv2d_layout(attrs, inputs, tinfos, relay) + + def expected(): + x = relay.var("x", shape=(1, 32, 56, 56)) + w = relay.var("w", shape=(32, 1, 3, 3)) + x = relay.layout_transform(x, "NCHW", "NCHW8c") + w = relay.layout_transform(w, "OIHW", "OIHW1i8o") + y = relay.nn.contrib_depthwise_conv2d_nchwc(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), + groups=32, data_layout="NCHW8c", kernel_layout="OIHW1i8o", + out_layout="NCHW8c") + y = relay.layout_transform(y, "NCHW8c", "NCHW") + y = relay.Function(free_vars(y), y) + return y + + a = before() + a = infer_type(a) + a = canonicalize_ops(a) + a = infer_type(a) + a = alter_op_layout(a) + a = infer_type(a) + + b = expected() + b = infer_type(b) + + assert(alpha_equal(a, b)) + + + if __name__ == "__main__": test_alter_op() @@ -572,3 +612,4 @@ def expected(): test_alter_layout_nchw_upsamping_op() test_alter_layout_strided_slice() test_alter_layout_prelu() + test_alter_layout_depthwise_conv2d() From 262a9c85edad2b53e523a1ce5e9abf3df4a0f794 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Fri, 31 May 2019 15:25:45 -0700 Subject: [PATCH 4/8] Fix test --- .../python/relay/test_pass_alter_op_layout.py | 82 +++++++++---------- 1 file changed, 41 insertions(+), 41 deletions(-) diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index 1a0f29cb6e80..7d022ba25570 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. """Test alter op layout pass""" +import tvm from tvm import relay from tvm.relay.op import register_alter_op_layout @@ -513,37 +514,30 @@ def expected(): assert alpha_equal(a, b), "Actual = \n" + str(a) -def test_alter_layout_prelu(): - """Test PRelu operator""" +def test_alter_layout_depthwise_conv2d(): + """Test depthwise_conv2d operator""" def before(): - x = relay.var("x", shape=(1, 64, 56, 56)) - weight = relay.var("weight") - alpha = relay.var("alpha", relay.IncompleteType()) - y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) - y = relay.nn.prelu(y, alpha) + x = relay.var("x", shape=(1, 32, 56, 56)) + w = relay.var("w", shape=(32, 1, 3, 3)) + y = relay.nn.conv2d(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), groups=32) y = relay.Function(free_vars(y), y) return y + import topi @register_alter_op_layout("nn.conv2d", level=110) def alter_conv2d(attrs, inputs, tinfos): - data, weight = inputs - new_attrs = dict(attrs) - new_attrs['data_layout'] = 'NCHW16c' - return relay.nn.conv2d(data, weight, **new_attrs) + with tvm.target.create("llvm"): + return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, relay) def expected(): - x = relay.var("x", shape=(1, 64, 56, 56)) - w = relay.var("weight") - alpha = relay.var("alpha", relay.IncompleteType()) - - y = relay.layout_transform(x, "NCHW", "NCHW16c") - y = relay.nn.conv2d(y, w, - channels=64, - kernel_size=(3, 3), - padding=(1, 1), - data_layout="NCHW16c") - y = relay.layout_transform(y, "NCHW16c", "NCHW") - y = relay.nn.prelu(y, alpha) + x = relay.var("x", shape=(1, 32, 56, 56)) + w = relay.var("w", shape=(32, 1, 3, 3)) + x = relay.layout_transform(x, "NCHW", "NCHW8c") + w = relay.layout_transform(w, "OIHW", "OIHW1i8o") + y = relay.nn.contrib_depthwise_conv2d_nchwc(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), + groups=32, data_layout="NCHW8c", kernel_layout="OIHW1i8o", + out_layout="NCHW8c") + y = relay.layout_transform(y, "NCHW8c", "NCHW") y = relay.Function(free_vars(y), y) return y @@ -559,29 +553,37 @@ def expected(): assert(alpha_equal(a, b)) -def test_alter_layout_depthwise_conv2d(): - """Test depthwise_conv2d operator""" +def test_alter_layout_prelu(): + """Test PRelu operator""" def before(): - x = relay.var("x", shape=(1, 32, 56, 56)) - w = relay.var("w", shape=(32, 1, 3, 3)) - y = relay.nn.conv2d(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), groups=32) + x = relay.var("x", shape=(1, 64, 56, 56)) + weight = relay.var("weight") + alpha = relay.var("alpha", relay.IncompleteType()) + y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)) + y = relay.nn.prelu(y, alpha) y = relay.Function(free_vars(y), y) return y - import topi @register_alter_op_layout("nn.conv2d", level=111) def alter_conv2d(attrs, inputs, tinfos): - return topi.x86.conv2d._alter_conv2d_layout(attrs, inputs, tinfos, relay) + data, weight = inputs + new_attrs = dict(attrs) + new_attrs['data_layout'] = 'NCHW16c' + return relay.nn.conv2d(data, weight, **new_attrs) def expected(): - x = relay.var("x", shape=(1, 32, 56, 56)) - w = relay.var("w", shape=(32, 1, 3, 3)) - x = relay.layout_transform(x, "NCHW", "NCHW8c") - w = relay.layout_transform(w, "OIHW", "OIHW1i8o") - y = relay.nn.contrib_depthwise_conv2d_nchwc(x, w, padding=(1, 1), channels=32, kernel_size=(3, 3), - groups=32, data_layout="NCHW8c", kernel_layout="OIHW1i8o", - out_layout="NCHW8c") - y = relay.layout_transform(y, "NCHW8c", "NCHW") + x = relay.var("x", shape=(1, 64, 56, 56)) + w = relay.var("weight") + alpha = relay.var("alpha", relay.IncompleteType()) + + y = relay.layout_transform(x, "NCHW", "NCHW16c") + y = relay.nn.conv2d(y, w, + channels=64, + kernel_size=(3, 3), + padding=(1, 1), + data_layout="NCHW16c") + y = relay.layout_transform(y, "NCHW16c", "NCHW") + y = relay.nn.prelu(y, alpha) y = relay.Function(free_vars(y), y) return y @@ -598,8 +600,6 @@ def expected(): assert(alpha_equal(a, b)) - - if __name__ == "__main__": test_alter_op() test_alter_return_none() @@ -611,5 +611,5 @@ def expected(): test_alter_layout_concatenate() test_alter_layout_nchw_upsamping_op() test_alter_layout_strided_slice() - test_alter_layout_prelu() test_alter_layout_depthwise_conv2d() + test_alter_layout_prelu() From 395408f187098699a4b09f3608d52290b65822a2 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Mon, 3 Jun 2019 09:52:44 -0700 Subject: [PATCH 5/8] Assert kernel layout --- topi/python/topi/x86/conv2d.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py index 8a709b8af1e3..316be188e348 100644 --- a/topi/python/topi/x86/conv2d.py +++ b/topi/python/topi/x86/conv2d.py @@ -415,8 +415,6 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F): dtype = data.dtype out_dtype = dtype if out_dtype in ("same", "") else out_dtype - kshape = get_const_tuple(kernel.shape) - is_depthwise = groups == kshape[0] and kshape[1] == 1 # only optimize for NCHW if layout != 'NCHW': @@ -424,6 +422,10 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F): if groups != 1 and not is_depthwise: return None + assert attrs["kernel_layout"] == "OIHW" + kshape = get_const_tuple(kernel.shape) + is_depthwise = groups == kshape[0] and kshape[1] == 1 + dispatch_ctx = autotvm.task.DispatchContext.current target = tvm.target.current_target() # query schedule and fallback if necessary From 39c622ec66d4da48d8c0d4b522e46a383d3e4d36 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Mon, 3 Jun 2019 13:30:46 -0700 Subject: [PATCH 6/8] Minor fix --- topi/python/topi/x86/conv2d.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py index 316be188e348..5775d80f9242 100644 --- a/topi/python/topi/x86/conv2d.py +++ b/topi/python/topi/x86/conv2d.py @@ -419,13 +419,14 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F): # only optimize for NCHW if layout != 'NCHW': return None - if groups != 1 and not is_depthwise: - return None assert attrs["kernel_layout"] == "OIHW" kshape = get_const_tuple(kernel.shape) is_depthwise = groups == kshape[0] and kshape[1] == 1 + if groups != 1 and not is_depthwise: + return None + dispatch_ctx = autotvm.task.DispatchContext.current target = tvm.target.current_target() # query schedule and fallback if necessary From f88585f45dfa4d6792852816dca7ee2a6bc19de8 Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Wed, 5 Jun 2019 15:02:50 -0700 Subject: [PATCH 7/8] Add get_shape function --- topi/python/topi/util.py | 39 +++++++++++++++++++++++++++++ topi/python/topi/x86/conv2d.py | 5 ++-- topi/tests/python/test_topi_util.py | 35 ++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 3 deletions(-) create mode 100644 topi/tests/python/test_topi_util.py diff --git a/topi/python/topi/util.py b/topi/python/topi/util.py index f648245c6bb7..260170738847 100644 --- a/topi/python/topi/util.py +++ b/topi/python/topi/util.py @@ -20,6 +20,7 @@ from numbers import Integral import tvm +from tvm.api import layout, bijective_layout from . import tag def traverse_inline(s, final_op, callback): @@ -297,3 +298,41 @@ def get_max_power2_factor(n, max_value=None): x *= 2 n /= 2 return x + + +def get_shape(src_shape, src_layout, dst_layout): + """Given a source shape, a source layout and a destination layout, infer + the destination shape. + + Parameter + --------- + src_shape : tuple of int or IntImm + Source shape + + src_layout : str or Layout + Source layout + + dst_layout : str or Layout + Destination layout + + Returns + ------- + dst_shape : tuple of int + Destination shape + """ + if src_layout == dst_layout: + return get_const_tuple(src_shape) + + if isinstance(src_layout, str): + src_layout = layout(src_layout) + if isinstance(dst_layout, str): + dst_layout = layout(dst_layout) + + assert len(src_layout) == len(dst_layout), \ + "Incompatible layout %s vs %s" % (src_layout, dst_layout) + + layout_mapping = bijective_layout(src_layout, dst_layout) + dst_indices = layout_mapping.forward_index( + tvm.convert([i for i in range(len(src_layout))])) + + return get_const_tuple(tuple([src_shape[i.value] for i in dst_indices])) diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py index 5775d80f9242..7cf821f9898d 100644 --- a/topi/python/topi/x86/conv2d.py +++ b/topi/python/topi/x86/conv2d.py @@ -26,7 +26,7 @@ from tvm.autotvm.task import get_config from .. import generic, tag from .. import nn -from ..util import get_const_tuple +from ..util import get_const_tuple, get_shape from ..nn.conv2d import conv2d, conv2d_NCHWc, \ conv2d_alter_layout, conv2d_infer_layout, _get_workload as _get_conv2d_workload from ..nn.depthwise_conv2d import _get_workload as _get_depthwise_conv2d_workload @@ -420,8 +420,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F): if layout != 'NCHW': return None - assert attrs["kernel_layout"] == "OIHW" - kshape = get_const_tuple(kernel.shape) + kshape = get_shape(kernel.shape, attrs["kernel_layout"], "OIHW") is_depthwise = groups == kshape[0] and kshape[1] == 1 if groups != 1 and not is_depthwise: diff --git a/topi/tests/python/test_topi_util.py b/topi/tests/python/test_topi_util.py new file mode 100644 index 000000000000..534b6993d411 --- /dev/null +++ b/topi/tests/python/test_topi_util.py @@ -0,0 +1,35 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Test code for util""" + +import topi + + +def verify_get_shape(src_shape, src_layout, dst_layout, expect_shape): + dst_shape = topi.util.get_shape(src_shape, src_layout, dst_layout) + assert dst_shape == expect_shape, \ + "Shape mismatch: expecting %s but got %s" % (expect_shape, dst_shape) + + +def test_get_shape(): + verify_get_shape((1, 3, 224, 224), "NCHW", "NCHW", (1, 3, 224, 224)) + verify_get_shape((1, 3, 224, 224), "NCHW", "NHWC", (1, 224, 224, 3)) + verify_get_shape((3, 2, 32, 48, 16), "NCHW16c", "NC16cWH", (3, 2, 16, 48, 32)) + verify_get_shape((2, 3, 32, 32, 16, 8), "OIHW16i8o", "HWO8oI16i", (32, 32, 2, 8, 3, 16)) + +if __name__ == "__main__": + test_get_shape() \ No newline at end of file From b81bf16ecaa745aa7cea3afe78315501b769850e Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Wed, 5 Jun 2019 15:21:59 -0700 Subject: [PATCH 8/8] Minor change --- topi/python/topi/x86/conv2d.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py index 7cf821f9898d..82d1caaccfc2 100644 --- a/topi/python/topi/x86/conv2d.py +++ b/topi/python/topi/x86/conv2d.py @@ -416,13 +416,13 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F): dtype = data.dtype out_dtype = dtype if out_dtype in ("same", "") else out_dtype - # only optimize for NCHW - if layout != 'NCHW': - return None - kshape = get_shape(kernel.shape, attrs["kernel_layout"], "OIHW") is_depthwise = groups == kshape[0] and kshape[1] == 1 + # only optimize for NCHW + if layout != 'NCHW' or attrs["kernel_layout"] != "OIHW": + return None + if groups != 1 and not is_depthwise: return None