[mlir][linalg] Implement LinalgGroupedConvolutionOpInterface to unify grouped convs

[srcarroll]

This patch aims to unify all grouped convs (and any subset) into a single linalg.grouped_conv_nd operation. This op takes any number of spatial dimensions but they are restricted to being contiguous. The order of batch, channel, group, and spatial dims for operands are specified by the layouts attribute. Since depthwise convolution is a special case of grouped when num groups is equal to the number of input channels, then depthwise convs are also covered by this op. Similarly, normal convolution is when num groups is 1, so those should also be covered.

This single op can replace all of the following

linalg.conv_1d_ncw_fcw (linalg::Conv1DNcwFcwOp)
linalg.conv_1d_nwc_wcf (linalg::Conv1DNwcWcfOp)
linalg.conv_1d (linalg::Conv1DOp)
linalg.conv_2d_nchw_fchw (linalg::Conv2DNchwFchwOp)
linalg.conv_2d_ngchw_fgchw (linalg::Conv2DNgchwFgchwOp)
linalg.conv_2d_ngchw_gfchw (linalg::Conv2DNgchwGfchwOp)
linalg.conv_2d_nhwc_fhwc (linalg::Conv2DNhwcFhwcOp)
linalg.conv_2d_nhwc_hwcf (linalg::Conv2DNhwcHwcfOp)
linalg.conv_2d_nhwgc_gfhwc (linalg::Conv2DNhwgcGfhwcOp)
linalg.conv_2d (linalg::Conv2DOp)
linalg.conv_3d_ncdhw_fcdhw (linalg::Conv3DNcdhwFcdhwOp)
linalg.conv_3d_ndhwc_dhwcf (linalg::Conv3DNdhwcDhwcfOp)

The LinalgGroupedConvolutionOpInterface would also allow us to represent the quantized versions of all these ops. However, it is not clear yet whether it is more desirable to cram quantization semantics into the same op. If it is, then there is no real need for the interface.

[ivangarcia44]

Hi @srcarroll , I was looking to add support for 3D grouped convolution and found this PR. It seems it includes the generalization of the grouped convolution to ND. Do you have plans to merge this PR? Please let me know if you need help finishing this change. Thanks,
Ivan

[srcarroll]

Hi @ivangarcia44, sure I'd be willing to bring it back to life if there's some need for it. I'd definitely hate for another conv op to be added and is why I started this in the first place. I just lost momentum and time to finish it. It might help to get some initial comments on what I have. I'll try to find some time this weekend to update it as it's pretty old by now

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[srcarroll]

there's probably no need for this interface if GroupedConvNDOp is already enough to cover this category without additional ops

[srcarroll]

i remember why i made an interface now. it was to also support implementing quantized versions of grouped conv, since at the time i didn't think we should cram that into the same op definition.

[ivangarcia44]

Hi @ivangarcia44, sure I'd be willing to bring it back to life if there's some need for it. I'd definitely hate for another conv op to be added and is why I started this in the first place. I just lost momentum and time to finish it. It might help to get some initial comments on what I have. I'll try to find some time this weekend to update it as it's pretty old by now

Thank you for taking this back. This could help in networks like the one in this paper for MRI's/CT scans: "A 3D Grouped Convolutional Network Fused with Conditional Random Field and Its Application in Image Multi‑target Fine Segmentation". Please let me know if any help is needed. Thanks

[srcarroll]

There has been an ongoing effort by @rengolin and others to thwart the op explosion in linalg dialect. See here where I had a brief discussion about it and here for the larger scoped discussion. So it would help most to get initial feedback from those involved. I believe @MaheshRavishankar and @banach-space are the relevant people to get eyes on this. Since I started this about a year ago, I'm not really sure how this fits into what they plan, so I dont want to get in the way. But if they feel this is in the right direction, then I can start working on this more

[srcarroll]

@MaheshRavishankar and @banach-space I went ahead and added you as reviewers but I wouldn't say it's necessarily ready for review. Just want an initial impression to see if it's worth continuing on this or not. thanks.

I know there was a PR going on here but seems stale. I haven't looked closely at it to see how much overlap there is. One thing to note is that this PR doesn't stop at 3 spatial dims. Although I've forgot a lot of details about convolutions, I think grouped convolutions also contain normal ones by setting the number of groups to 1

[srcarroll]

This single op can replace all of the following
linalg.conv_1d_ncw_fcw (linalg::Conv1DNcwFcwOp)
linalg.conv_1d_nwc_wcf (linalg::Conv1DNwcWcfOp)
linalg.conv_1d (linalg::Conv1DOp)
linalg.conv_2d_nchw_fchw (linalg::Conv2DNchwFchwOp)
linalg.conv_2d_ngchw_fgchw (linalg::Conv2DNgchwFgchwOp)
linalg.conv_2d_ngchw_gfchw (linalg::Conv2DNgchwGfchwOp)
linalg.conv_2d_nhwc_fhwc (linalg::Conv2DNhwcFhwcOp)
linalg.conv_2d_nhwc_hwcf (linalg::Conv2DNhwcHwcfOp)
linalg.conv_2d_nhwgc_gfhwc (linalg::Conv2DNhwgcGfhwcOp)
linalg.conv_2d (linalg::Conv2DOp)
linalg.conv_3d_ncdhw_fcdhw (linalg::Conv3DNcdhwFcdhwOp)
linalg.conv_3d_ndhwc_dhwcf (linalg::Conv3DNdhwcDhwcfOp)

I should say has the potential to replace. it already does replace all the grouped variants. but replacing normal convs, for example, might require more finessing. Like ignoring singleton dims due to group sizes of 1 for example