Group Convolution

Group convolution

@373 = gpu::pooling[mode=average,padding={0, 0, 0, 0},padding_mode=0,stride={1, 1},lengths={3, 3},dilations={1, 1},ceil_mode=0,lp_order=2,dyn_global=0](@371,@372) -> float_type, {1, 768, 17, 17}, {221952, 289, 17, 1}
@399 = gpu::code_object[code_object=7832,symbol_name=mlir_slice_convolution_add_relu,global=14592,local=64,](@397,@378,@390,@398) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@402 = gpu::code_object[code_object=7952,symbol_name=mlir_convolution_add_relu,global=14592,local=64,](@400,@396,@383,@401) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@404 = gpu::code_object[code_object=7352,symbol_name=mlir_convolution,global=14592,local=64,](@402,@381,@403) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@405 = load[offset=221952,end=443904](@1) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@406 = gpu::code_object[code_object=7480,symbol_name=mlir_convolution,global=14592,local=64,](@399,@388,@405) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@407 = load[offset=665856,end=887808](@1) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@408 = gpu::code_object[code_object=6392,symbol_name=mlir_convolution,global=14592,local=64,](@373,@380,@407) -> float_type, {1, 192, 17, 17}, {55488, 289, 17, 1}
@409 = broadcast[axis=1,out_lens={1, 576, 17, 17}](@379) -> float_type, {1, 576, 17, 17}, {0, 1, 0, 0}
@410 = load[offset=2848384,end=3514240](@1) -> float_type, {1, 576, 17, 17}, {166464, 289, 17, 1}
@411 = gpu::code_object[code_object=9496,symbol_name=concat_add_relu_kernel,global=55488,local=1024,](@406,@404,@408,@409,@410) -> float_type, {1, 576, 17, 17}, {166464, 289, 17, 1}

We could fuse these convolutions together using a group convolution. We would first reshape the inputs to make the group a seperate dimension so that:

{1, 768, 17, 17} => {1, 1, 768, 17, 17}
{1, 192, 17, 17} => {1, 1, 192, 17, 17}

Then concat across this new dimension. Then do another reshape to collapse it with the channels and then do group convolution.

key there here is to ensure input channels are the same

eg)https://paperswithcode.com/method/grouped-convolution