[microNPU][ETHOSU] Add offloading to the NPU the nn.avg_pool2d operator with a stride > 3 #14861
Conversation
|
cc @neildhickey, @ekalda, @ilyag-grovety, @Aleksei-grovety @arina-grovety |
arina-grovety
force-pushed
the
ethosu-avgpool-high-strides
branch
from
e9e9f3c to
fd1fcc0
Compare
There was a problem hiding this comment.
Thanks @sergio-grovety, looks good in general, nice documentation! Left some suggestions and questions
| ifm: (1, 8, 8, _), strides: (8, 8), pool_shape: (8, 8) | ||
| ifm: (1, 25, 5, _), strides: (25, 5), pool_shape: (25, 5) | ||
| """ | ||
| if self.pooling_type != "AVG": |
There was a problem hiding this comment.
Why can't we use the same trick for max pool as well?
There was a problem hiding this comment.
Hi @ekalda, exactly, this approach can also be applied to MAX_POOL, done
| return False | ||
| if [self.ifm.shape[1], self.ifm.shape[2]] != list(self.pool_shape): | ||
| return False | ||
| if list(self.strides) != list(self.pool_shape): |
There was a problem hiding this comment.
Why is this condition necessary? To me it sounds like this patch supports a case where the kernel is the same shape as ifm, so the stride value doesn't matter since we can't move the kernel anyway?
There was a problem hiding this comment.
Yes, you're right, now I see redundancy as well.
| @@ -650,12 +650,33 @@ def is_valid(self): | |||
| """ | |||
| This function checks whether AvgPool2D has compatible attributes with the NPU | |||
| """ | |||
|
|
|||
| def check_extra_strides(): | |||
There was a problem hiding this comment.
Maybe we can come up with a more self explanatory name for this function? It currently sounds as if some operators have more than one set of strides. Maybe something like check_same_ifm_and_kernel_shape or legalize_large_stride?
| @@ -875,6 +875,91 @@ def verify(ext_func): | |||
| verify(mod["tvmgen_default_ethos_u_main_0"]) | |||
|
|
|||
|
|
|||
| def test_tflite_pool2d_extra_strides_legalize(): | |||
There was a problem hiding this comment.
This legalization test is fine, it checks whether the pooling matching the specific conditions gets partitioned for microNPU. I think it would be good to also test that it gets correctly legalized to ethosu.pooling op (i.e. the strides will be correctly replaced with [1, 1]).
| # is [0, 0, 0, 0], the application of the pooling kernel will be done only once, | ||
| # which will give us the desired output | ||
| if params.pooling_type == "AVG" and (params.strides[0] > 3 or params.strides[1] > 3): | ||
| new_strides = [1, 1] |
There was a problem hiding this comment.
Maybe add strides variable to leave one return in the function.
… Extended the test_tflite_pool2d_same_ifm_and_kernel_shape_legalize().
sergio-grovety
force-pushed
the
ethosu-avgpool-high-strides
branch
from
fd1fcc0 to
e960179
Compare
|
Thanks @arina-grovety and @sergio-grovety, this is now merged! |
…or with a stride > 3 (apache#14861) The nn.avg_pool2d operator with a stride size greater than 3 in any of the spatial dimensions is rewritten as ethosu avg_pool with strides=[1,1] if the case satisfies the additional conditions (no AvgPool2D padding, spatial dimensions of ifm and shape of pooling are equal). --------- Co-authored-by: Sergey Smirnov <[email protected]> Co-authored-by: arina-grovety <> Co-authored-by: Arina Naumova (grovety.com) <[email protected]> Co-authored-by: Arina <[email protected]>
The nn.avg_pool2d operator with a stride size greater than 3 in any of the spatial dimensions is rewritten as ethosu avg_pool with strides=[1,1] if the case satisfies the additional conditions (no AvgPool2D padding, spatial dimensions of ifm and shape of pooling are equal).