Distance function

Since we operates on vectors, the most obvious use of SIMD instructions is to vectorize the distance procedure calculate_distance.

A decent compiler is able to auto-vectorize such a simple function; GCC and clang require flag -ffast-math to enable this. Here is the assembly code generated by GCC 7.3.1:

.L7:
        movaps  (%r10,%rax), %xmm1
        addl    $1, %r8d
        movups  (%rcx,%rax), %xmm3
        addq    $16, %rax
        cmpl    %r11d, %r8d
        subps   %xmm3, %xmm1
        mulps   %xmm1, %xmm1
        addps   %xmm1, %xmm2
        jb      .L7

Below is a hand-coded version, without processing vectors' tail.

float sse_dist_vector(size_t n, float* a, float* b) {

    __m128 ret = _mm_setzero_ps();

    for (size_t i=0; i < n; i += 4) {

        const __m128 ai = _mm_loadu_ps(a + i);
        const __m128 bi = _mm_loadu_ps(b + i);

        const __m128 d  = _mm_sub_ps(ai, bi);
        const __m128 d2 = _mm_mul_ps(d, d);

        ret = _mm_add_ps(ret, d2);
    }

    float v[4];
    _mm_storeu_ps(v, ret);

    return v[0] + v[1] + v[2] + v[3];
}

Multiple distances at once

In the inner loop the expression centroid[i] is constant. We can use this fact and calculate multiple distances at once, minimizing number of memory fetches from centroids array and also utilizing all SIMD resources.

So the algorithm becomes:

for (int i=0; i < K; i++) {
    calculate_multiple_distances(centroid[i], samples, distances);
}

And actual implementation of calculate_multiple_distances, which processes four vectors at once:

void sse_dist_vector_many(size_t n, float const_vector[], size_t k, float* in[], float out[]) {

    for (size_t j=0; j < k; j += 4) {

        float* in0 = in[j + 0]; // Note: factoring out these pointer gave pretty
        float* in1 = in[j + 1]; //       nice boost, as GCC 7.3.1 couldn't handle
        float* in2 = in[j + 2]; //       well complex addresses like `in[j + 0][i]`
        float* in3 = in[j + 3];

        __m128 out0 = _mm_setzero_ps();
        __m128 out1 = _mm_setzero_ps();
        __m128 out2 = _mm_setzero_ps();
        __m128 out3 = _mm_setzero_ps();

        for (size_t i=0; i < n; i += 4) {
            const __m128 vc = _mm_loadu_ps(const_vector + i);

            const __m128 x0 = _mm_loadu_ps(in0 + i);
            const __m128 x1 = _mm_loadu_ps(in1 + i);
            const __m128 x2 = _mm_loadu_ps(in2 + i);
            const __m128 x3 = _mm_loadu_ps(in3 + i);

            const __m128 d0   = _mm_sub_ps(x0, vc);
            const __m128 d0_2 = _mm_mul_ps(d0, d0);
            const __m128 d1   = _mm_sub_ps(x1, vc);
            const __m128 d1_2 = _mm_mul_ps(d1, d1);
            const __m128 d2   = _mm_sub_ps(x2, vc);
            const __m128 d2_2 = _mm_mul_ps(d2, d2);
            const __m128 d3   = _mm_sub_ps(x3, vc);
            const __m128 d3_2 = _mm_mul_ps(d3, d3);

            out0 = _mm_add_ps(out0, d0_2);
            out1 = _mm_add_ps(out1, d1_2);
            out2 = _mm_add_ps(out2, d2_2);
            out3 = _mm_add_ps(out3, d3_2);
        }

        float v[4];

        _mm_storeu_ps(v, out0);
        out[j + 0] = v[0] + v[1] + v[2] + v[3];

        _mm_storeu_ps(v, out1);
        out[j + 1] = v[0] + v[1] + v[2] + v[3];

        _mm_storeu_ps(v, out2);
        out[j + 2] = v[0] + v[1] + v[2] + v[3];

        _mm_storeu_ps(v, out3);
        out[j + 3] = v[0] + v[1] + v[2] + v[3];
    }
}

While in a naive algorithm we fetch numbers from the centroids table K * N times, here we do K * N/4, just 25%.