Improved consistency of CPU dispatching code

pgvector · Apr 18, 2024 · 4e093f9 · 4e093f9
1 parent fb3c964
commit 4e093f9
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 19 deletions.
diff --git a/src/bitutils.h b/src/bitutils.h
@@ -2,19 +2,24 @@
 #define BITUTILS_H
 
 /* Only enable for more recent compilers */
+/* TODO Move to better place */
 #if defined(__x86_64__) && defined(__GNUC__) && __GNUC__ >= 8
-#define BIT_DISPATCH
+#define USE_DISPATCH
 #elif defined(__x86_64__) && defined(__clang_major__) && __clang_major__ >= 7
-#define BIT_DISPATCH
+#define USE_DISPATCH
 #elif defined(_M_AMD64) && defined(_MSC_VER) && _MSC_VER >= 1920
-#define BIT_DISPATCH
+#define USE_DISPATCH
 #endif
 
 /* target_clones requires glibc */
-#if defined(BIT_DISPATCH) && defined(__gnu_linux__)
+#if defined(USE_DISPATCH) && defined(__gnu_linux__)
 #define USE_TARGET_CLONES
 #endif
 
+#if defined(USE_DISPATCH)
+#define BIT_DISPATCH
+#endif
+
 extern uint64 (*BitHammingDistance) (uint32 bytes, unsigned char *ax, unsigned char *bx, uint64 distance);
 extern double (*BitJaccardDistance) (uint32 bytes, unsigned char *ax, unsigned char *bx, uint64 ab, uint64 aa, uint64 bb);
 

diff --git a/src/halfvec.h b/src/halfvec.h
@@ -5,10 +5,11 @@
 
 #include <float.h>
 
+#include "bitutils.h"
 #include "fmgr.h"
 #include "vector.h"
 
-#if defined(__x86_64__) || defined(_M_AMD64)
+#if defined(USE_DISPATCH)
 #define HALFVEC_DISPATCH
 #endif
 

diff --git a/src/vector.c b/src/vector.c
@@ -34,18 +34,10 @@
 #define STATE_DIMS(x) (ARR_DIMS(x)[0] - 1)
 #define CreateStateDatums(dim) palloc(sizeof(Datum) * (dim + 1))
 
-/* target_clones requires glibc */
-#if defined(__gnu_linux__) && defined(__has_attribute)
-/* Use separate line for portability */
-#if __has_attribute(target_clones)
-#define HAVE_TARGET_CLONES
-#endif
-#endif
-
-#if defined(__x86_64__) && defined(HAVE_TARGET_CLONES) && !defined(__FMA__)
-#define VECTOR_DISPATCH __attribute__((target_clones("default", "fma")))
+#if defined(USE_TARGET_CLONES) && !defined(__FMA__)
+#define VECTOR_TARGET_CLONES __attribute__((target_clones("default", "fma")))
 #else
-#define VECTOR_DISPATCH
+#define VECTOR_TARGET_CLONES
 #endif
 
 PG_MODULE_MAGIC;
@@ -573,7 +565,7 @@ halfvec_to_vector(PG_FUNCTION_ARGS)
 	PG_RETURN_POINTER(result);
 }
 
-VECTOR_DISPATCH static float
+VECTOR_TARGET_CLONES static float
 VectorL2SquaredDistance(int dim, float *ax, float *bx)
 {
 	float		distance = 0.0;
@@ -620,7 +612,7 @@ vector_l2_squared_distance(PG_FUNCTION_ARGS)
 	PG_RETURN_FLOAT8((double) VectorL2SquaredDistance(a->dim, a->x, b->x));
 }
 
-VECTOR_DISPATCH static float
+VECTOR_TARGET_CLONES static float
 VectorInnerProduct(int dim, float *ax, float *bx)
 {
 	float		distance = 0.0;
@@ -662,7 +654,7 @@ vector_negative_inner_product(PG_FUNCTION_ARGS)
 	PG_RETURN_FLOAT8((double) -VectorInnerProduct(a->dim, a->x, b->x));
 }
 
-VECTOR_DISPATCH static double
+VECTOR_TARGET_CLONES static double
 VectorCosineSimilarity(int dim, float *ax, float *bx)
 {
 	float		similarity = 0.0;