nearest neighbor

1nn.c

#include <stdio.h>
#include <string.h>
#include <time.h>
#include <math.h>
#include <stdlib.h>
#include <stdint.h>
#include <immintrin.h>
#include "sys/param.h"


float euclidean_distance_sq(float *vecs, int veclen, int a, int b) {
	float d = 0;
	for (int i = 0; i < veclen; i++) {
		float t = vecs[veclen * a + i] - vecs[veclen * b + i];
		d += t * t;
	}
	return d;
}


double euclidean_distance_sq_vec(double *vecs, int veclen, int a, int b) {
	__m256d d = _mm256_set1_pd(0.0);
	for (int i = 0; i < veclen; i+=4) {
		__m256d aa = _mm256_load_pd(vecs + veclen * a + i);
		__m256d bb = _mm256_load_pd(vecs + veclen * b + i);
		__m256d t = _mm256_sub_pd(aa, bb);
		d = _mm256_add_pd(d, _mm256_mul_pd(t, t));
	}

	__m256d sum = _mm256_hadd_pd(d, d);
	return ((double*)&sum)[0] + ((double*)&sum)[2];
}


float euclidean_distance_sq_vec_fl(float *vecs, int veclen, int a, int b) {
	__m256 d = _mm256_set1_ps(0.0);
	for (int i = 0; i < veclen; i+=8) {
		__m256 aa = _mm256_load_ps(vecs + veclen * a + i);
		__m256 bb = _mm256_load_ps(vecs + veclen * b + i);
		__m256 t = _mm256_sub_ps(aa, bb);
		d = _mm256_add_ps(d, _mm256_mul_ps(t, t));
	}

	d = _mm256_hadd_ps(d, d);
	d = _mm256_hadd_ps(d, d);
	return ((float*)&d)[0] + ((float*)&d)[4];
}


float euclidean_distance_sq_ssse(float *vecs, int veclen, int a, int b) {
	__m128 d = _mm_set1_ps(0.0f);
	for (int i = 0; i < veclen; i+=4) {
		__m128 aa = _mm_load_ps(vecs + veclen * a + i);
		__m128 bb = _mm_load_ps(vecs + veclen * b + i);
		__m128 t = _mm_sub_ps(aa, bb);
		d = _mm_add_ps(d, _mm_mul_ps(t, t));
	}

	__m128 sum = _mm_hadd_ps(d, d);
	return ((float*)&sum)[0] + ((float*)&sum)[1];
}


int nn(float *vecs, int veccnt, int veclen, int idx) {
	int nearest = -1;
	float dist = INFINITY;
	for (int i = 0; i < veccnt; i++) {
		float d = euclidean_distance_sq(vecs, veclen, idx, i);
		if (i != idx && d < dist) {
			nearest = i;
			dist = d;
		}
	}
	return nearest;
}


struct idxdist {
	int idx;
	float dist;
};


int main(int argc, char *argv[]) {

	int veccnt = atoi(argv[1]);
	int veclen = atoi(argv[2]);

	if (veclen % 8 != 0) {
		printf("veclen not divisible by 8\n");
		return 1;
	}

	printf("count: %d, vector length: %d\n", veccnt, veclen);

	float *vecs;
	posix_memalign((void**)&vecs, 32, sizeof(float) * veccnt * veclen);

	struct idxdist *idxdists = malloc(sizeof(struct idxdist) * veccnt);

	srand(47);
	for (int i = 0; i < veccnt * veclen; i++) {
		vecs[i] = ((float) rand() / (float) RAND_MAX);
	}

	for (int i = 0; i < veccnt; i++) {
		idxdists[i].idx = -1;
		idxdists[i].dist = INFINITY;
	}

	/*
	for (int i = 0; i < veccnt; i++) {
		idxdists[i].idx = nn(vecs, veccnt, veclen, i);
	}
	*/

	
	int tile = 64;

	#pragma omp parallel for
	for (int ti = 0; ti < veccnt; ti += tile) {

		for (int tj = 0; tj < veccnt; tj += tile) {
			int maxi = MIN(ti+tile, veccnt);
			for (int i = ti; i < maxi; i++) { // current vec idx
				int maxj = MIN(tj+tile, veccnt);
				for (int j = tj; j < maxj; j++) {

					float d = euclidean_distance_sq_vec_fl(vecs, veclen, i, j);
					if (j != i && d < idxdists[i].dist) {
						idxdists[i].idx = j;
						idxdists[i].dist = d;
					}

				}
			}
		}
	}

	for (int i = 0; i < MIN(32, veccnt); i++) {
		printf("%d ", idxdists[i].idx);
	}
	printf("\n");

}