Fast "index of minimum" for distances to points on single thread

Question

The problem in simple words: I have three arrays of double values (x-, y- and z-coordinates) and need to find the index of the point which has minimum distance to a reference points. I do not need the actual value of this distance. If the minimum distance is not unique, any index where its value occurs may be returned.

I will be using this function for determining the squared distance between points:

double distSqr(double xa, double ya, double za,
               double xb, double yb, double zb)
{
    double X = xa - xb;
    double Y = xa - xb;
    double Z = xa - xb;

    return X * X + Y * Y + Z * Z;
}

Taking the square root in the final step will not be necessary.

First of all, consider a "naive" implementation where you simply compare the "current" value to all other values, exchanging the "current" value if a smaller value is found:

// Take the three arrays of x-, y-, and z-coordinates filled with arbitrary values.
#define SIZE 20

double cx[SIZE];
double cy[SIZE];
double cz[SIZE];

int min_naive(double x, double y, double z)
{
    double minDist = distSqr(x, y, z, cx[0], cy[0], cz[0]);
    int minIndex = 0;

    for (int i = 1; i < SIZE; i++)
    {
        double d = distSqr(x, y, z, cx[i], cy[i], cz[i]);
        if (d < minDist)
        {
            minDist = d;
            minIndex = i;
        }
    }

    return minIndex;
}

// Finally, determin the index, e.g.
int minInd = min_naive(0, 0, 0);

This does not seem like an optimal solution to me, especially since it might not be fast enough if the loop has to run over all distances where only one comparison happens at each step. Obviously, sooner or later all distances have to be calculated. Next, I came up with a better solution that first calculates all distances and then compares them in a loop in such a way, that each loop halves the number of possible distances left. I hoped that the compiler might auto-vectorize this function, instead of using a simple loop over all points, but I have not checked that yet.

This is the code I came up with:

void cmpmv(int lower, int upper, int *inds, double *dists)
{
    double a = dists[lower];
    double b = dists[upper];

    if (a > b)
    {
        inds[lower] = inds[upper];
        dists[lower] = b;
    }
}

int min_better(double kax, double kay, double kaz)
{
    double dists[SIZE];
    int inds[SIZE];
    for (int i = 0; i < SIZE; i++)
    {
        dists[i] = distSqr(kax, kay, kaz, nncx[i], nncy[i], nncz[i]);
        inds[i] = i;
    }

    int div, mod;
    int s = SIZE;
    while (s > 1)
    {
        div = s / 2;
        mod = s % 2;

        for (int i = 0; i < div; i++)
            cmpmv(i, i + div, inds, dists);
        if (mod == 1)
            cmpmv(0, s - 1, inds, dists);

        s = div;
    }

    return inds[0];
}

(Note that this must differentiate the cases where SIZE and later s are not a multiple of 2.)

Are there maybe even better and especially faster ways to implement this?

There are two restrictions:

1. I only have one thread available.
2. I have to use arrays/pointers.

Answer 1

The major compilers did not really auto-vectorize this, but it can be done manually. For example with AVX, we could do something like (not tested)

int indexOfMin(double pt_x, double pt_y, double pt_z, int n)
{
    __m256d ptx = _mm256_set1_pd(pt_x);
    __m256d pty = _mm256_set1_pd(pt_y);
    __m256d ptz = _mm256_set1_pd(pt_z);
    __m256d xdif = _mm256_sub_pd(ptx, _mm256_loadu_pd(cx));
    __m256d ydif = _mm256_sub_pd(pty, _mm256_loadu_pd(cy));
    __m256d zdif = _mm256_sub_pd(ptz, _mm256_loadu_pd(cz));
    __m256d min_dist = _mm256_add_pd(_mm256_add_pd(_mm256_mul_pd(xdif, xdif), 
                                                   _mm256_mul_pd(ydif, ydif)), 
                                                   _mm256_mul_pd(zdif, zdif));
    __m128i min_index = _mm_set_epi32(3, 2, 1, 0);
    __m128i index = min_index;
    __m256d dist;
    for (int i = 4; i < n; i += 4) {
        xdif = _mm256_sub_pd(ptx, _mm256_load_pd(cx + i));
        ydif = _mm256_sub_pd(pty, _mm256_load_pd(cy + i));
        zdif = _mm256_sub_pd(ptz, _mm256_load_pd(cz + i));
        dist = _mm256_add_pd(_mm256_add_pd(_mm256_mul_pd(xdif, xdif), 
                                           _mm256_mul_pd(ydif, ydif)), 
                                           _mm256_mul_pd(zdif, zdif));
        index = _mm_add_epi32(index, _mm_set1_epi32(4));
        __m256 mask256 = _mm256_castpd_ps(_mm256_cmp_pd(dist, min_dist, _CMP_LT_OS));
        // mask256 has the masks as 4 x int64, but we need 4 x int32
        // there's no nice 'pack' to do it, but shufps can extract
        // the relevant floats, and then we can reinterpret as integers
        // mask256 = * D * C * B * A (* is an ignored float)
        __m128 maskL = _mm256_castps256_ps128(mask256);   // * B * A
        __m128 maskH = _mm256_extractf128_ps(mask256, 1); // * D * C
        __m128 maskps = _mm_shuffle_ps(maskL, maskH, _MM_SHUFFLE(2, 0, 2, 0)); // D C B A
        __m128i mask = _mm_castps_si128(maskps);
        min_dist = _mm256_min_pd(min_dist, dist);
        // if the mask is set (this distance is LT the old minimum) then take the current index
        // otherwise keep the old index
        min_index = _mm_blendv_epi8(min_index, index, mask);
    }

    double mdist[4];
    _mm256_storeu_pd(mdist, min_dist);
    uint32_t mindex[4];
    _mm_storeu_si128((__m128i*)mindex, min_index);
    double closest = mdist[0];
    int closest_i = mindex[0];
    for (int i = 1; i < 4; i++) {
        if (mdist[i] < closest) {
            closest = mdist[i];
            closest_i = mindex[i];
        }
    }
    return closest_i;
}

The relevant header to include is <immintrin.h> and to compile you would need to enable AVX with -mavx (GCC, Clang) or /arch:AVX (MSVC).

Most of the code is just subtracting the values from the given coordinate, squaring the difference, and summing the squares. That's not very interesting to discuss, though it plays a significant part in making the code fast. Finding the minimum is more interesting, and is what prevented auto-vectorization. The approach I used is comparing the distance (obviously that was going to be part of it) which results in a bit-mask that is all set where the comparison is true, and then I use that to blend between the "index of best-so-far" and the current index, to conditionally replace values without branching.

Because AVX was targeted instead of AVX2, the simpler approach of using an __m256i for the indexes could not be used. That would have removed the need to extract/shuffle the mask, as it would have already been the right size. With AVX2 there are no 256 bit wide integer operations (well, mostly) so it would not be possible to add 4 to the indexes. It would be possible to do a 256 bit blend by using a floating point type blend, then the blend mask is easy but it just pushes the problem to incrementing the indexes.

Finally at the end there is a small loop to select the best index among the 4 candidates.

The array size must be a multiple of 4, it would not be hard to remove that requirement. 32-byte alignment of the arrays is not required but would be better.

By the way this is related to an SSE2 version that I did a couple of years ago, which used single precision floats.