[Mlir-commits] [mlir] [mlir][linalg] Add TransposeConv2D Pass (PR #68567)

[Jack Frankland]

FranklandJack wrote:

> > > Thanks for the contribution! My main comment would be that this pass is a bit limited ATM:
> > > > Convert conv_2d_nhwc_fhwc to conv_2d_nhwc_hwcf
> > > 
> > > 
> > > Why not support conversions in both directions?
> > 
> > 
> > Sure, although I'm not 100% clear what "both directions" means here. Should the pass also run on `conv_2d_nhwc_hwcf` and convert it to `fwhc`?
> 
> Yes, that's what I had in mind. And in general, it would be helpful if you expanded a bit on the use cases for this in the summary (and how would this be used). Otherwise it feels a bit like some arbitrary transformation. I am sure it isn't, but it's just not clear from the summary. More specifically, is there a hardware that works better with one type over the other? And some frameworks tend to generate the "wrong" type? Just curious, but I suspect other people might be too.
> 
> Also, apologies for the delay responding and thanks for the updates :)

No worries about the delay! Thanks for the thorough reivew! 🤗 

Okay so the motivation behind this is that we made [some changes in the TOSA to LinAlg lowering to allow a more direct mapping](https://github.com/llvm/llvm-project/pull/68304). Previously `tosa.conv2d` (which has the `{f,h,w,c}` ordering) was lowered to linalg's `linalg.conv_2d_nhwc_hwcf`. In order to reorder the filter channels a tranpose was materialized (essentially what this pass does was previously part of TOSA->LinAlg lowering). (By the way as an aside, the TOSA documentatino refers to the filter as "weights" so I guess this is where I got that from 🤦)

We felt that a more direct mapping was appropriate as there is a 1-1 correspondence between `tosa.conv2d` and LinAlg's `linalg.conv_2d_nhwc_fhwc` and that it should be the responsibility of further transformations to introduce the tranpose if the pipeline in question wants that ordering.

`{h,w,c,f}` can be a useful ordering for things like img2col where the contiguous `{c,f}` filter dimensions allow you to employ a GEMM (which can be useful if you support optimized routines / BLAS libraries). On the other hand we now also [support an img2col type transformation for the `{f,h,w,c}` filter ordering](https://github.com/llvm/llvm-project/pull/68708), which could in theory allow better cache coherency as all memory accesses happen in a contiguous way.

So I guess in summary, this pass allows anyone who wants to restore the old behaviour to do so but in a way such that they know they are materializing the extra transposition. I hope this helps, sorry for rambling. I agree this is quite a specific transformation but I'm not sure how else to add it without a whole pass.

I guess I'm still a little confused about the utility of converting `{h,w,c,f}` to `{f,w,h,c}` here when the purpose is to sort of align on the `{h,w,c,f}` ordering for the filters, perhaps we could add an option or something to define which way you want to go?

https://github.com/llvm/llvm-project/pull/68567