Added SIMD version of L2 distance

ankane · ankane · commit 925aa4e048f0 · 2024-04-07T20:22:19.000-07:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -123,6 +123,6 @@ jobs:
       - uses: ankane/setup-postgres-valgrind@v1
         with:
           postgres-version: 16
-      - run: make
+      - run: make OPTFLAGS=""
       - run: sudo --preserve-env=PG_CONFIG make install
       - run: make installcheck
diff --git a/src/halfvec.c b/src/halfvec.c
@@ -810,13 +810,42 @@ l2_distance_squared_internal(HalfVector * a, HalfVector * b)
 	half	   *bx = b->x;
 	float		distance = 0.0;
 
+#if defined(F16C_SUPPORT) && defined(__FMA__)
+	int			i;
+	float		s[8];
+	int			count = (a->dim / 8) * 8;
+	__m256		dist = _mm256_setzero_ps();
+
+	for (i = 0; i < count; i += 8)
+	{
+		__m128i		axi = _mm_loadu_si128((__m128i *) (ax + i));
+		__m128i		bxi = _mm_loadu_si128((__m128i *) (bx + i));
+		__m256		axs = _mm256_cvtph_ps(axi);
+		__m256		bxs = _mm256_cvtph_ps(bxi);
+		__m256		diff = _mm256_sub_ps(axs, bxs);
+
+		dist = _mm256_fmadd_ps(diff, diff, dist);
+	}
+
+	_mm256_store_ps(s, dist);
+
+	distance = s[0] + s[1] + s[2] + s[3] + s[4] + s[5] + s[6] + s[7];
+
+	for (; i < a->dim; i++)
+	{
+		float		diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
+
+		distance += diff * diff;
+	}
+#else
 	/* Auto-vectorized */
 	for (int i = 0; i < a->dim; i++)
 	{
 		float		diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
 
 		distance += diff * diff;
 	}
+#endif
 
 	return (double) distance;
 }
