liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-12 15:46:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fixed CPU dispatching check for halfvec distance functions

Browse Source
This commit is contained in:
Andrew Kane
2024-04-17 18:27:27 -07:00
parent e2d8907180
commit fb6af03361
1 changed files with 18 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
src/halfutils.c
View File

@@ -192,13 +192,20 @@ HalfvecCosineSimilarityF16cFma(int dim, half * ax, half * bx)
#endif #endif
#ifdef HALFVEC_DISPATCH #ifdef HALFVEC_DISPATCH
#define CPU_FEATURE_FMA (1 << 12) #define CPU_FEATURE_FMA (1 << 12)
#define CPU_FEATURE_F16C (1 << 29) #define CPU_FEATURE_OSXSAVE (1 << 27)
#define CPU_FEATURE_AVX (1 << 28)
#define CPU_FEATURE_F16C (1 << 29)
static bool #ifdef _MSC_VER
#define TARGET_XSAVE
#else
#define TARGET_XSAVE __attribute__((target("xsave")))
#endif
TARGET_XSAVE static bool
SupportsCpuFeature(unsigned int feature) SupportsCpuFeature(unsigned int feature)
{ {
/* TODO Fix check */
unsigned int exx[4] = {0, 0, 0, 0}; unsigned int exx[4] = {0, 0, 0, 0};
#if defined(HAVE__GET_CPUID) #if defined(HAVE__GET_CPUID)
@@ -207,6 +214,12 @@ SupportsCpuFeature(unsigned int feature)
__cpuid(exx, 1); __cpuid(exx, 1);
#endif #endif
if ((exx[2] & CPU_FEATURE_OSXSAVE) != CPU_FEATURE_OSXSAVE)
return false;
if ((_xgetbv(0) & 6) != 6)
return false;
return (exx[2] & feature) == feature; return (exx[2] & feature) == feature;
} }
#endif #endif
@@ -223,7 +236,7 @@ HalfvecInit(void)
HalfvecCosineSimilarity = HalfvecCosineSimilarityDefault; HalfvecCosineSimilarity = HalfvecCosineSimilarityDefault;
#ifdef HALFVEC_DISPATCH #ifdef HALFVEC_DISPATCH
if (SupportsCpuFeature(CPU_FEATURE_FMA | CPU_FEATURE_F16C)) if (SupportsCpuFeature(CPU_FEATURE_AVX | CPU_FEATURE_F16C | CPU_FEATURE_FMA))
{ {
HalfvecL2SquaredDistance = HalfvecL2SquaredDistanceF16cFma; HalfvecL2SquaredDistance = HalfvecL2SquaredDistanceF16cFma;
HalfvecInnerProduct = HalfvecInnerProductF16cFma; HalfvecInnerProduct = HalfvecInnerProductF16cFma;