Optimization of 3D Direct Convolution Implementation in C

问题

For my project, I've written a naive C implementation of direct 3D convolution with periodic padding on the input. Unfortunately, since I'm new to C, the performance isn't so good... here's the code:

int mod(int a, int b)
{
    // calculate mod to get the correct index with periodic padding
    int r = a % b;
    return r < 0 ? r + b : r;
}
void convolve3D(const double *image, const double *kernel, const int imageDimX, const int imageDimY, const int imageDimZ, const int stencilDimX, const int stencilDimY, const int stencilDimZ, double *result)
{
    int imageSize = imageDimX * imageDimY * imageDimZ;
    int kernelSize = kernelDimX * kernelDimY * kernelDimZ;

    int i, j, k, l, m, n;
    int kernelCenterX = (kernelDimX - 1) / 2;
    int kernelCenterY = (kernelDimY - 1) / 2;
    int kernelCenterZ = (kernelDimZ - 1) / 2;
    int xShift,yShift,zShift;
    int outIndex, outI, outJ, outK;
    int imageIndex = 0, kernelIndex = 0;
    
    // Loop through each voxel
    for (k = 0; k < imageDimZ; k++){
        for ( j = 0; j < imageDimY; j++) {
            for ( i = 0; i < imageDimX; i++) {
                stencilIndex = 0;
                // for each voxel, loop through each kernel coefficient
                for (n = 0; n < kernelDimZ; n++){
                    for ( m = 0; m < kernelDimY; m++) {
                        for ( l = 0; l < kernelDimX; l++) {
                            // find the index of the corresponding voxel in the output image
                            xShift = l - kernelCenterX;
                            yShift = m - kernelCenterY;
                            zShift = n - kernelCenterZ;

                            outI = mod ((i - xShift), imageDimX);
                            outJ = mod ((j - yShift), imageDimY);
                            outK = mod ((k - zShift), imageDimZ);
                            
                            outIndex = outK * imageDimX * imageDimY + outJ * imageDimX + outI;

                            // calculate and add
                            result[outIndex] += stencil[stencilIndex]* image[imageIndex];
                            stencilIndex++;
                        }
                    }
                } 
                imageIndex ++;
            }
        }
    } 
}

by convention, all the matrices (image, kernel, result) are stored in column-major fashion, and that's why I loop through them in such way so they are closer in memory (heard this would help).

I know the implementation is very naive, but since it's written in C, I was hoping the performance would be good, but instead it's a little disappointing. I tested it with image of size 100^3 and kernel of size 10^3 (Total ~1GFLOPS if only count the multiplication and addition), and it took ~7s, which I believe is way below the capability of a typical CPU.

If possible, could you guys help me optimize this routine? I'm open to anything that could help, with just a few things if you could consider:

The problem I'm working with could be big (e.g. image of size 200 by 200 by 200 with kernel of size 50 by 50 by 50 or even larger). I understand that one way of optimizing this is by converting this problem into a matrix multiplication problem and use the blas GEMM routine, but I'm afraid memory could not hold such a big matrix
Due to the nature of the problem, I would prefer direct convolution instead of FFTConvolve, since my model is developed with direct convolution in mind, and my impression of FFT convolve is that it gives slightly different result than direct convolve especially for rapidly changing image, a discrepancy I'm trying to avoid. That said, I'm in no way an expert in this. so if you have a great implementation based on FFTconvolve and/or my impression on FFT convolve is totally biased, I would really appreciate if you could help me out.
The input images are assumed to be periodic, so periodic padding is necessary
I understand that utilizing blas/SIMD or other lower level ways would definitely help a lot here. but since I'm a newbie here I dont't really know where to start... I would really appreciate if you help pointing me to the right direction if you have experience in these libraries,

Thanks a lot for your help, and please let me know if you need more info about the nature of the problem

回答1:

As a first step, replace your mod ((i - xShift), imageDimX) with something like this:

inline int clamp( int x, int size )
{
    if( x < 0 ) return x + size;
    if( x >= size ) return x - size;
    return x;
}

These branches are very predictable because they yield same results for very large count of consecutive elements. Integer modulo is relatively slow.

Now, next step (ordered by cost/profit) is going to be parallelizing. If you have any modern C++ compiler, just enable OpenMP somewhere in project settings. After that you need 2 changes.

Decorate your very outer loop with something like this: #pragma omp parallel for schedule(guided)
Move your function-level variables within that loop. This also means you’ll have to compute initial imageIndex from your k, for each iteration.

Next option, rework your code so you only write each output value once. Compute the final value in your innermost 3 loops, reading from random locations from both image and kernel, and only write the result once. When you have that result[outIndex] += in the inner loop, CPU stalls waiting for the data from memory. When you accumulate in a variable that’s a register not memory, there’s no access latency.

SIMD is the most complicated optimization for that. But in short, you’ll need maximum width of the FMA your hardware has (if you have AVX and need double precision, that width is 4), and you’ll also need multiple independent accumulators for your 3 innermost loops, to avoid hitting the latency as opposed to saturating the throughput. Here’s my answer to much easier problem as an example what I mean.

来源：https://stackoverflow.com/questions/62603086/optimization-of-3d-direct-convolution-implementation-in-c

标签

optimization

simd

convolution

blas