From c621db94daebc3ec49a6fee75cc920ad55a36bef Mon Sep 17 00:00:00 2001 From: farshidsp Date: Wed, 16 Nov 2022 23:02:25 -0800 Subject: [PATCH 1/9] enable depthwise conv2d NHWC with HWIO kernel layout --- python/tvm/relay/op/strategy/hexagon.py | 4 ++-- python/tvm/relay/op/strategy/x86.py | 4 ++-- python/tvm/topi/nn/depthwise_conv2d.py | 23 +++++++++++++++++------ 3 files changed, 21 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/op/strategy/hexagon.py b/python/tvm/relay/op/strategy/hexagon.py index c1d64f2fe143..e4cc63f6737d 100644 --- a/python/tvm/relay/op/strategy/hexagon.py +++ b/python/tvm/relay/op/strategy/hexagon.py @@ -86,9 +86,9 @@ def conv2d_strategy_hexagon(attrs, inputs, out_type, target): name="depthwise_conv2d_nchw.hexagon", ) elif layout == "NHWC": - assert kernel_layout == "HWOI" strategy.add_implementation( - wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc), + wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, + need_kernel_layout=True), wrap_topi_schedule(topi.hexagon.schedule_depthwise_conv2d_nhwc), name="depthwise_conv2d_nhwc.hexagon", ) diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py index 10d7fbb3a926..c1f1d61dc579 100644 --- a/python/tvm/relay/op/strategy/x86.py +++ b/python/tvm/relay/op/strategy/x86.py @@ -228,13 +228,13 @@ def conv2d_strategy_cpu(attrs, inputs, out_type, target): assert _OIHWio_matcher.match(kernel_layout) # check if kernel is OIHWio return depthwise_conv2d_NCHWc_strategy_cpu(attrs, inputs, out_type, target) elif layout == "NHWC": - assert kernel_layout == "HWOI" if (not need_auto_scheduler_layout) and (not need_meta_schedule_layout): logger.warning( "depthwise_conv2d NHWC layout is not optimized for x86 with autotvm." ) strategy.add_implementation( - wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc), + wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, + need_kernel_layout=True), wrap_topi_schedule(topi.generic.schedule_depthwise_conv2d_nhwc), name="depthwise_conv2d_nhwc.generic", ) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index 48ffb8c6d9ff..4ac327cfb80f 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -19,6 +19,7 @@ from __future__ import absolute_import as _abs from collections import namedtuple import tvm +import numpy as np from tvm import te from .dilate import dilate @@ -211,7 +212,7 @@ def depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=No return Output -def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, out_dtype=None): +def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layout, out_dtype=None): """Depthwise convolution nhwc forward operator. Parameters @@ -252,9 +253,19 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, out_dtype=No dilation_h, dilation_w = dilation batch, in_height, in_width, in_channel = Input.shape + + dim = len(Input.shape) - 2 + # shape of dilated kernel - filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape - + if kernel_layout == "HWOI": + filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape + kernel_permutation_to = [0, 1] + list(range(2, dim + 2)) + elif kernel_layout == "HWIO": + filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape + kernel_permutation_to = [dim + 1, dim] + list(range(dim)) + + kernel_permutation_from = np.argsort(kernel_permutation_to) + dilated_kernel_h = (filter_height - 1) * dilation_h + 1 dilated_kernel_w = (filter_width - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( @@ -284,9 +295,9 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, out_dtype=No j * stride_w + dj * dilation_w, idxdiv(c, channel_multiplier), ].astype(out_dtype) - * Filter[ - di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier) - ].astype(out_dtype) + * Filter.__getitem__(tuple(np.array([ + di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)])[kernel_permutation_from]) + ).astype(out_dtype) ), axis=[di, dj], ), From fc2b365c8de075f1ab23cb919e21f5b19e41dfec Mon Sep 17 00:00:00 2001 From: farshidsp Date: Thu, 17 Nov 2022 01:39:29 -0800 Subject: [PATCH 2/9] lint fix --- python/tvm/relay/op/strategy/hexagon.py | 3 +-- python/tvm/relay/op/strategy/x86.py | 3 +-- python/tvm/topi/nn/depthwise_conv2d.py | 20 ++++++++++++-------- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/python/tvm/relay/op/strategy/hexagon.py b/python/tvm/relay/op/strategy/hexagon.py index e4cc63f6737d..f42503a1477c 100644 --- a/python/tvm/relay/op/strategy/hexagon.py +++ b/python/tvm/relay/op/strategy/hexagon.py @@ -87,8 +87,7 @@ def conv2d_strategy_hexagon(attrs, inputs, out_type, target): ) elif layout == "NHWC": strategy.add_implementation( - wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, - need_kernel_layout=True), + wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, need_kernel_layout=True), wrap_topi_schedule(topi.hexagon.schedule_depthwise_conv2d_nhwc), name="depthwise_conv2d_nhwc.hexagon", ) diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py index c1f1d61dc579..d9bedacb1678 100644 --- a/python/tvm/relay/op/strategy/x86.py +++ b/python/tvm/relay/op/strategy/x86.py @@ -233,8 +233,7 @@ def conv2d_strategy_cpu(attrs, inputs, out_type, target): "depthwise_conv2d NHWC layout is not optimized for x86 with autotvm." ) strategy.add_implementation( - wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, - need_kernel_layout=True), + wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, need_kernel_layout=True), wrap_topi_schedule(topi.generic.schedule_depthwise_conv2d_nhwc), name="depthwise_conv2d_nhwc.generic", ) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index 4ac327cfb80f..dc4404ba5c17 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -253,19 +253,19 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou dilation_h, dilation_w = dilation batch, in_height, in_width, in_channel = Input.shape - + dim = len(Input.shape) - 2 - + # shape of dilated kernel if kernel_layout == "HWOI": - filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape - kernel_permutation_to = [0, 1] + list(range(2, dim + 2)) + filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape + kernel_permutation_to = [0, 1] + list(range(2, dim + 2)) elif kernel_layout == "HWIO": filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape kernel_permutation_to = [dim + 1, dim] + list(range(dim)) - + kernel_permutation_from = np.argsort(kernel_permutation_to) - + dilated_kernel_h = (filter_height - 1) * dilation_h + 1 dilated_kernel_w = (filter_width - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( @@ -295,8 +295,12 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou j * stride_w + dj * dilation_w, idxdiv(c, channel_multiplier), ].astype(out_dtype) - * Filter.__getitem__(tuple(np.array([ - di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)])[kernel_permutation_from]) + * Filter.__getitem__( + tuple( + np.array( + [di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)] + )[kernel_permutation_from] + ) ).astype(out_dtype) ), axis=[di, dj], From 6cdb2d0ab231f67cea17fe953b4f959dd8424774 Mon Sep 17 00:00:00 2001 From: farshidsp Date: Wed, 30 Nov 2022 17:34:25 -0800 Subject: [PATCH 3/9] bug fix --- python/tvm/topi/nn/depthwise_conv2d.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index dc4404ba5c17..c13f46d491c6 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -262,10 +262,9 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou kernel_permutation_to = [0, 1] + list(range(2, dim + 2)) elif kernel_layout == "HWIO": filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape - kernel_permutation_to = [dim + 1, dim] + list(range(dim)) + kernel_permutation_to = list(range(dim)) + [dim + 1, dim] kernel_permutation_from = np.argsort(kernel_permutation_to) - dilated_kernel_h = (filter_height - 1) * dilation_h + 1 dilated_kernel_w = (filter_width - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( @@ -295,13 +294,13 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou j * stride_w + dj * dilation_w, idxdiv(c, channel_multiplier), ].astype(out_dtype) - * Filter.__getitem__( + * Filter[ tuple( np.array( [di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)] )[kernel_permutation_from] ) - ).astype(out_dtype) + ].astype(out_dtype) ), axis=[di, dj], ), From e582efb30268c82cde83d80129e8354daf9786d6 Mon Sep 17 00:00:00 2001 From: farshidsp Date: Mon, 5 Dec 2022 09:30:22 -0800 Subject: [PATCH 4/9] fix --- python/tvm/topi/nn/depthwise_conv2d.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index c13f46d491c6..a1134772dd1d 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -254,17 +254,16 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou batch, in_height, in_width, in_channel = Input.shape - dim = len(Input.shape) - 2 - # shape of dilated kernel if kernel_layout == "HWOI": filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape - kernel_permutation_to = [0, 1] + list(range(2, dim + 2)) + kernel_permutaiton = [0, 1, 2, 3] elif kernel_layout == "HWIO": filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape - kernel_permutation_to = list(range(dim)) + [dim + 1, dim] + kernel_permutaiton = [0, 1, 3, 2] + else: + raise ValueError(f"Unsupported kernel layout: {kernel_layout}") - kernel_permutation_from = np.argsort(kernel_permutation_to) dilated_kernel_h = (filter_height - 1) * dilation_h + 1 dilated_kernel_w = (filter_width - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( @@ -298,7 +297,7 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou tuple( np.array( [di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)] - )[kernel_permutation_from] + )[kernel_permutaiton] ) ].astype(out_dtype) ), From ce70659b8cc5fe540daf6a8dcef774d44e052618 Mon Sep 17 00:00:00 2001 From: farshidsp Date: Mon, 5 Dec 2022 09:31:32 -0800 Subject: [PATCH 5/9] fix --- python/tvm/topi/nn/depthwise_conv2d.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index a1134772dd1d..fc6a52247f95 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -257,10 +257,10 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou # shape of dilated kernel if kernel_layout == "HWOI": filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape - kernel_permutaiton = [0, 1, 2, 3] + kernel_permutation = [0, 1, 2, 3] elif kernel_layout == "HWIO": filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape - kernel_permutaiton = [0, 1, 3, 2] + kernel_permutation = [0, 1, 3, 2] else: raise ValueError(f"Unsupported kernel layout: {kernel_layout}") @@ -297,7 +297,7 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou tuple( np.array( [di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)] - )[kernel_permutaiton] + )[kernel_permutation] ) ].astype(out_dtype) ), From e34374b2fafe7f67893a984a2618c61d9779c2d2 Mon Sep 17 00:00:00 2001 From: farshidsp Date: Mon, 5 Dec 2022 15:05:41 -0800 Subject: [PATCH 6/9] fix --- python/tvm/topi/nn/depthwise_conv2d.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index fc6a52247f95..3ce2f7eb687d 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -255,14 +255,12 @@ def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layou batch, in_height, in_width, in_channel = Input.shape # shape of dilated kernel - if kernel_layout == "HWOI": - filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape - kernel_permutation = [0, 1, 2, 3] - elif kernel_layout == "HWIO": + if kernel_layout == "HWIO": filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape kernel_permutation = [0, 1, 3, 2] else: - raise ValueError(f"Unsupported kernel layout: {kernel_layout}") + filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape + kernel_permutation = [0, 1, 2, 3] dilated_kernel_h = (filter_height - 1) * dilation_h + 1 dilated_kernel_w = (filter_width - 1) * dilation_w + 1 From 0bf433f4a86d0acae660ae6fe967a08b6e408cd4 Mon Sep 17 00:00:00 2001 From: farshidsp Date: Tue, 6 Dec 2022 12:42:12 -0800 Subject: [PATCH 7/9] have the default kernel layout to HWOI --- python/tvm/topi/nn/depthwise_conv2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index 3ce2f7eb687d..1d6a2310beea 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -212,7 +212,7 @@ def depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=No return Output -def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layout, out_dtype=None): +def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layout="HWOI", out_dtype=None): """Depthwise convolution nhwc forward operator. Parameters From de68e6701a1916e4ceac3166910da1903eb090ad Mon Sep 17 00:00:00 2001 From: farshidsp Date: Tue, 6 Dec 2022 12:43:04 -0800 Subject: [PATCH 8/9] lint --- python/tvm/topi/nn/depthwise_conv2d.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py index 1d6a2310beea..7c446a23a813 100644 --- a/python/tvm/topi/nn/depthwise_conv2d.py +++ b/python/tvm/topi/nn/depthwise_conv2d.py @@ -212,7 +212,9 @@ def depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=No return Output -def depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layout="HWOI", out_dtype=None): +def depthwise_conv2d_nhwc( + Input, Filter, stride, padding, dilation, kernel_layout="HWOI", out_dtype=None +): """Depthwise convolution nhwc forward operator. Parameters From 5155ebd9c67f5575b69753b5f3a8cf6e76f64171 Mon Sep 17 00:00:00 2001 From: farshidsp Date: Mon, 12 Dec 2022 12:24:27 -0800 Subject: [PATCH 9/9] arm-cpu fix --- python/tvm/relay/op/strategy/arm_cpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/op/strategy/arm_cpu.py b/python/tvm/relay/op/strategy/arm_cpu.py index 261b979dedaf..c8d51bc23c82 100644 --- a/python/tvm/relay/op/strategy/arm_cpu.py +++ b/python/tvm/relay/op/strategy/arm_cpu.py @@ -318,7 +318,7 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type, target): else: logger.warning("depthwise_conv2d with layout NHWC is not optimized for arm cpu.") strategy.add_implementation( - wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc), + wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc, need_kernel_layout=True), wrap_topi_schedule(conv2d_generic.schedule_depthwise_conv2d_nhwc), name="depthwise_conv2d_nhwc.generic", )