Improved performance of binary_quantize function for halfvec

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.8.1 (unreleased)
 
 - Added support for Postgres 18
+- Improved performance of `binary_quantize` function
 
 ## 0.8.0 (2024-10-30)
 

src/halfvec.c

@@ -898,8 +898,21 @@ halfvec_binary_quantize(PG_FUNCTION_ARGS)
 	half	   *ax = a->x;
 	VarBit	   *result = InitBitVector(a->dim);
 	unsigned char *rx = VARBITS(result);
+	int			i = 0;
+	int			count = (a->dim / 8) * 8;
 
-	for (int i = 0; i < a->dim; i++)
+	/* Auto-vectorized */
+	for (; i < count; i += 8)
+	{
+		unsigned char result_byte = 0;
+
+		for (int j = 0; j < 8; j++)
+			result_byte |= (HalfToFloat4(ax[i + j]) > 0) << (7 - j);
+
+		rx[i / 8] = result_byte;
+	}
+
+	for (; i < a->dim; i++)
 		rx[i / 8] |= (HalfToFloat4(ax[i]) > 0) << (7 - (i % 8));
 
 	PG_RETURN_VARBIT_P(result);

src/vector.c

@@ -946,17 +946,20 @@ binary_quantize(PG_FUNCTION_ARGS)
 	float	   *ax = a->x;
 	VarBit	   *result = InitBitVector(a->dim);
 	unsigned char *rx = VARBITS(result);
-	int			i;
+	int			i = 0;
 	int			count = (a->dim / 8) * 8;
-	unsigned char result_byte;
 
-	for (i = 0; i < count; i += 8)
+	/* Auto-vectorized */
+	for (; i < count; i += 8)
 	{
-		result_byte = 0;
+		unsigned char result_byte = 0;
+
 		for (int j = 0; j < 8; j++)
 			result_byte |= (ax[i + j] > 0) << (7 - j);
+
 		rx[i / 8] = result_byte;
 	}
+
 	for (; i < a->dim; i++)
 		rx[i / 8] |= (ax[i] > 0) << (7 - (i % 8));