Improve NHWC depthwise convolution for AArch64

[giuseros]

High level description of this contribution

We created a default schedule (no auto-tuning or tensorization) named
depthwise_conv2d_nhwc which does a decent job at optimizing depthwise
for NHWC layouts (on AArch64 architectures).

The schedule lives in : topi/python/topi/arm_cpu/depthwise_conv2d.py
While we register the strategy in : python/tvm/relay/op/strategy/arm_cpu.py

This contribution is based on the following RFC: https://discuss.tvm.ai/t/rfc-improve-depthwise-convolution-for-nhwc-layouts-on-aarch64-targets/7360

[giuseros]

cc @u99127 @anijain2305 @FrozenGene

[giuseros]

Hi @FrozenGene ,
Before introducing tuning knobs, I wanted to do first an analysis to find the minimum set of tuning parameters to bring the best performance.

The aim is to reduce the tuning time. The point is that sometimes we are constrained by the number of registers available in AArch64, so trying out different splits might only increase the tuning time, without giving any benefit.

So the idea was to have a "default" schedule which mimics ACL implementation and then introduce (the minimal set of) tuning knobs + tensorization to speed things up.

What do you think? If you want to add tuning knobs in this PR, I will try to do the tuning analysis today

[FrozenGene]

ACL implementation

Hi @giuseros Thanks for the work. I fully understand your purpose and smoothy development path. As this schedule will be the default NHWC depthwise convolution, my opinion is we should try to achieve a good performance as far as we could achieve. Notably I don't mean we mush achieve like ACL ultimate performance then we could merge, optimization is not one-shot deal. But here I think we could enable auto tvm to help us to achieve better performance. I think it is worthy introducing into this pr.

• This schedule will be applied for arm32 and arm64 both, we shouldn't only consider arm64. So auto tvm (split) could help us to avoid this issue.

• Tuning knob of compute_at (especially data_pad) could help us solve parallel-compute-locality issue (we can not assume we only run kernel only in one single core). see more detail: http://people.csail.mit.edu/jrk/halide-pldi13.pdf Figure 2

I agree we should reduce tuning knob and improve tuning time experience, but if it could help us improve performance, I think we should introduce it in, otherwise we could avoid it.

[giuseros]

Hi @FrozenGene ,

Let me thank you for the review and the pointers (the Halide paper is quite interesting)!

I tried most of your suggestions and all I got was a tiny improvement on mobilenetv2 but a significant increase in tuning time.

Since this looks like it will take more time for more analysis I would prefer if we could take this PR in as is and I can follow up with further improvements in the future.

[giuseros]

Hi @tqchen ,

Is this an issue with my changes or with the CI? It seems to point to an import in a particular file, but I am not able to see anything wrong with that import.

[FrozenGene]

Hi @FrozenGene ,

Let me thank you for the review and the pointers (the Halide paper is quite interesting)!

I tried most of your suggestions and all I got was a tiny improvement on mobilenetv2 but a significant increase in tuning time.

Since this looks like it will take more time for more analysis I would prefer if we could take this PR in as is and I can follow up with further improvements in the future.

I agree. We could limit our tuning knob if we find it does no effect.

[FrozenGene]

Hi @tqchen ,

Is this an issue with my changes or with the CI? It seems to point to an import in a particular file, but I am not able to see anything wrong with that import.

Seems not related with your change. Others meet this CI error too.

[giuseros]

Hi @FrozenGene ,
Final changes to this PR.

• Introducing the compute_at knobs forces us to use xgb_knob tuner which does not support locate_cache annotations (so I switched to a custom knob).
• I also noticed that we were not legalizing depthwsie, which means we were running pooling, reductions, etc, in order to compute the offset contribution. Legalizing depthwise gives a boost of 2x (for quantized).

If this won't be approved by tonight, I will turn into a draft to pick up after I come back from holidays (i.e., 15 days).

[FrozenGene]

LGTM

[anijain2305]

LGTM. Just one minor comment.

[giuseros]

I probably enabled a CUDA test that was making the CI hang. I reverted the test hoping that this was the issue

[giuseros]

@FrozenGene , @anijain2305 ,
This is my last commit before holidays. I enabled only the arm_cpu tests (everything passes locally). The case with dilation>1 will be untested for now (as it is for CUDA). I hope this will make CI happy. If not, I will pick this up when I am back.

[giuseros]

Hi @FrozenGene , @anijain2305 ,
This PR finally passed the CI. Would it be possible to merge it?

Thanks!

[FrozenGene]

Thanks @giuseros Merged now