Decompose aten.channel_shuffle op (#4243)

[ivangarcia44]

Support for the channel shuffle operator is added by torch dialect level decomposition (similar to the pixel_shuffle operation).

The decomposition is based on this specification:
https://docs.pytorch.org/docs/stable/generated/torch.nn.ChannelShuffle.html
and implementation:
aten/src/ATen/native/ChanelShuffle.cpp
https://github.com/pytorch/pytorch/blob/23491519d288dedb2a54cfad5fef7fcb2ad8eade/aten/src/ATen/native/ChanelShuffle.cpp#L4

Note that the operator consists of an expansion, expanded channel dimensions permute, and contraction of channel dimensions back to the original size. For example, for an input array of shape 1x8x4x4 with a group size of 2 would generate the MLIR linalg code below.

module {
func.func @channel_shuffle(%arg0: !torch.vtensor<[1, 8, 4, 4], f32>) -> !torch.vtensor<[1, 8, 4, 4], f32> {
%c0 = torch.constant.int 0
%c1 = torch.constant.int 1
%c2 = torch.constant.int 2
%c3 = torch.constant.int 3
%c4 = torch.constant.int 4
%dims = torch.prim.ListConstruct %c0, %c2, %c1, %c3, %c4 : (!torch.int, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list

%reshaped = torch.prims.split_dim %arg0, %c1, %c2 : !torch.vtensor<[1, 8, 4, 4], f32>, !torch.int, !torch.int -> !torch.vtensor<[1, 4, 2, 4, 4], f32>

%permuted = torch.aten.permute %reshaped, %dims : !torch.vtensor<[1, 4, 2, 4, 4], f32>, !torch.list -> !torch.vtensor<[1, 2, 4, 4, 4], f32>

%collapsed = torch.prims.collapse %permuted, %c1, %c2 : !torch.vtensor<[1, 2, 4, 4, 4], f32>, !torch.int, !torch.int -> !torch.vtensor<[1, 8, 4, 4], f32>

return %collapsed : !torch.vtensor<[1, 8, 4, 4], f32>

}
}
References:
PyTorch ChannelShuffle definition:
https://docs.pytorch.org/docs/stable/generated/torch.nn.ChannelShuffle.html
ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices (2017):
https://arxiv.org/pdf/1707.01083
A Lightweight Dendritic ShuffleNet for Medical Image Classification (2025)
https://www.jstage.jst.go.jp/article/transinf/advpub/0/advpub_2024EDP7059/_pdf
PyTorch implementation:
aten/src/ATen/native/ChanelShuffle.cpp
https://github.com/pytorch/pytorch/blob/23491519d288dedb2a54cfad5fef7fcb2ad8eade/aten/src/ATen/native/ChanelShuffle.cpp#L4

Resolves #4243

[ivangarcia44]

Hi all, this is a reminder to please provide feedback to this PR to add support of the channel shuffle operation in torch to linalg lowering. Thank you

@newling @silvasean @rsuderman @zjgarvey @penguin-wwy @rafaelubalmw @sahas3 @vinitdeodhar @alaa-ali @dixinzhou @ramiro050 @qedawkins @vivekkhandelwal1

[sahas3]

LGTM with a minor nit.

Please wait for other reviewers to review before merging. Thanks!

[alaa-ali]

LGTM.

[ivangarcia44]

Hi all, this is a reminder to please provide feedback to this PR to add support of the channel shuffle operation in torch to linalg lowering. Thank you

@newling @silvasean @rsuderman @zjgarvey @penguin-wwy @rafaelubalmw @sahas3 @vinitdeodhar @alaa-ali @dixinzhou @ramiro050 @qedawkins @vivekkhandelwal1

[ivangarcia44]

Hi all, this is a reminder to please provide feedback to this PR to add support of the channel shuffle operation in torch to linalg lowering. Thank you

@newling @silvasean @rsuderman @zjgarvey @penguin-wwy @rafaelubalmw @sahas3 @vinitdeodhar @alaa-ali @dixinzhou @ramiro050 @qedawkins @vivekkhandelwal1 @stellaraccident

[stellaraccident]

This looks like it is thorough work and review. I apologize for the long review latency.

[stellaraccident]

Please let me know if you need this submitted for you.

[ivangarcia44]

Please let me know if you need this submitted for you.

@sahas3 will do the merge. Thank you for your help!