Add non-unit groups to aten::convolution

[gpetters94]

Adds support for group sizes > 1 for at::convolution (and by extension the ops that decompose into it).

[ramiro050]

Looking good! I have a few initial comments

[gpetters94]

Okay, most recent commit is confirmed working with my new op in LLVM. Once that lands and we can push torch-mlir's LLVM version up this should be good to go.

[gpetters94]

Now that the LLVM op has landed, here's the updated PR using the new logic. I spoke with Mahesh about adding a depthwise case and my concerns that it won't be possible with dynamic tensors because I won't be able to do the necessary calculation at compile time. He suggested that the optimization to depthwise convolution could be done on the backend (i.e. IREE) side. If you think I should add a specific case for depthwise static tensors then I can do that, otherwise this is complete.

[silvasean]

Now that the LLVM op has landed, here's the updated PR using the new logic. I spoke with Mahesh about adding a depthwise case and my concerns that it won't be possible with dynamic tensors because I won't be able to do the necessary calculation at compile time. He suggested that the optimization to depthwise convolution could be done on the backend (i.e. IREE) side. If you think I should add a specific case for depthwise static tensors then I can do that, otherwise this is complete.

I think we should support the special case when enough static info is available (the whole shape doesn't have to be static for us to detect it). It is a very important pattern in practice and we don't want all backends replicating that code.

[gpetters94]

Now that the LLVM op has landed, here's the updated PR using the new logic. I spoke with Mahesh about adding a depthwise case and my concerns that it won't be possible with dynamic tensors because I won't be able to do the necessary calculation at compile time. He suggested that the optimization to depthwise convolution could be done on the backend (i.e. IREE) side. If you think I should add a specific case for depthwise static tensors then I can do that, otherwise this is complete.

I think we should support the special case when enough static info is available (the whole shape doesn't have to be static for us to detect it). It is a very important pattern in practice and we don't want all backends replicating that code.

Okay, I've put in a PR at LLVM for the new op here: https://reviews.llvm.org/D128575. It's only for the case where the multiplier is 1, which covers Mobilenet. I can add the non-unit multiplier case if you think it's relevant, but for now I'm going to check for multiplier of 1 and fall through to the grouped op otherwise.

[gpetters94]

Testing my patch with other models (i.e. not Mobilenet) exposed an issue in the ordering of the op I wrote in upstream. This latest commit fixes the issue, although it does require a change to the upstream op which I have put up at https://reviews.llvm.org/D128880.

[gpetters94]

Both LLVM patches have landed, now this is just waiting for an LLVM bump before it's good to go.