Efficient 4D Convolution

[dCubelic]

Hi,

what would be the best way to schedule 4d convolution (NHWC), for kernels of size Nx3x3xC

this is the algorithm I am using:

auto paddedInput{ Halide::BoundaryConditions::constant_exterior( input, 0, 0, inputSizes[ 0 ], 0, inputSizes[ 1 ], 0, inputSizes[ 2 ], 0, inputSizes[ 3 ] ) };

auto heightPadding{ ( outputSize( 2 ) - 1 ) * stride_h + kernel.dim( 3 ).max() + 1 - inputSizes[ 2 ] };
auto widthPadding { ( outputSize( 1 ) - 1 ) * stride_w + kernel.dim( 2 ).max() + 1 - inputSizes[ 1 ] };

auto padTop { heightPadding / 2 };
auto padLeft{ widthPadding  / 2 };

Halide::RDom r( 0, kernel.dim( 1 ).max() + 1, 0, kernel.dim( 2 ).max() + 1, 0, kernel.dim( 3 ).max() + 1 );

output( c, w, h, n ) = 0.f;
output( c, w, h, n ) += paddedInput( r.x, w * stride_w - padLeft + r.y, h * stride_h - padTop + r.z, 0 ) * kernel( c, r.x, r.y, r.z );

currently the best results I could get are with this schedule I modified from the Adams2019 autoscheduler:

output.update(0)
    .split(c, c, ci, 8, Halide::TailStrategy::GuardWithIf)
    .parallel( h )
    .vectorize(ci)
    .reorder({ci, r.y, r.z, c, r.x, w, h, n})
    .unroll( ci );
output
    .split(c, c, ci, 8, Halide::TailStrategy::RoundUp)
    .parallel( h )
    .vectorize(ci)
    .compute_root()
    .reorder({ci, c, w, h, n});

I am sure it could be improved even more. Do you have any suggestions?

[abadams]

A good schedule for a conv layer can be quite complex. Take a look at https://github.com/halide/Halide/tree/main/apps/conv_layer