From 60b4bb2ad14d5c5b394d8421d9cdfaf6c6d87240 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 8 Apr 2024 10:00:34 -0700 Subject: [PATCH] Moved halfvec distance functions to separate file [skip ci] --- Makefile | 2 +- Makefile.win | 2 +- src/halfutils.c | 99 +++++++++++++++++++++++++++++++++++++++++++++++ src/halfutils.h | 9 +++++ src/halfvec.c | 100 +++--------------------------------------------- 5 files changed, 115 insertions(+), 97 deletions(-) create mode 100644 src/halfutils.c create mode 100644 src/halfutils.h diff --git a/Makefile b/Makefile index a7be0ef..d665f2e 100644 --- a/Makefile +++ b/Makefile @@ -3,7 +3,7 @@ EXTVERSION = 0.6.2 MODULE_big = vector DATA = $(wildcard sql/*--*.sql) -OBJS = src/bitvector.o src/halfvec.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/sparsevec.o src/vector.o +OBJS = src/bitvector.o src/halfutils.o src/halfvec.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/sparsevec.o src/vector.o HEADERS = src/halfvec.h src/sparsevec.h src/vector.h TESTS = $(wildcard test/sql/*.sql) diff --git a/Makefile.win b/Makefile.win index 48fd71b..fd598a8 100644 --- a/Makefile.win +++ b/Makefile.win @@ -1,7 +1,7 @@ EXTENSION = vector EXTVERSION = 0.6.2 -OBJS = src\bitvector.obj src\halfvec.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\sparsevec.obj src\vector.obj +OBJS = src\bitvector.obj src\halfutils.obj src\halfvec.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\sparsevec.obj src\vector.obj HEADERS = src\halfvec.h src\sparsevec.h src\vector.h REGRESS = bit_functions btree cast copy halfvec_functions halfvec_input hnsw_bit_hamming hnsw_bit_jaccard hnsw_halfvec_cosine hnsw_halfvec_ip hnsw_halfvec_l2 hnsw_options hnsw_sparsevec_cosine hnsw_sparsevec_ip hnsw_sparsevec_l2 hnsw_unlogged hnsw_vector_cosine hnsw_vector_ip hnsw_vector_l2 ivfflat_options ivfflat_unlogged ivfflat_vector_cosine ivfflat_vector_ip ivfflat_vector_l2 sparsevec_functions sparsevec_input vector_functions vector_input diff --git a/src/halfutils.c b/src/halfutils.c new file mode 100644 index 0000000..e961e05 --- /dev/null +++ b/src/halfutils.c @@ -0,0 +1,99 @@ +#include "postgres.h" + +#include "halfutils.h" +#include "halfvec.h" + +#ifdef F16C_SUPPORT +#include +#endif + +/* + * Get the L2 squared distance between half vectors + */ +double +HalfvecL2DistanceSquared(HalfVector * a, HalfVector * b) +{ + half *ax = a->x; + 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_storeu_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; +} + +/* + * Get the inner product of two half vectors + */ +double +HalfvecInnerProduct(HalfVector * a, HalfVector * b) +{ + half *ax = a->x; + 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); + + dist = _mm256_fmadd_ps(axs, bxs, dist); + } + + _mm256_storeu_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++) + distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]); +#else + /* Auto-vectorized */ + for (int i = 0; i < a->dim; i++) + distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]); +#endif + + return (double) distance; +} diff --git a/src/halfutils.h b/src/halfutils.h new file mode 100644 index 0000000..75562c8 --- /dev/null +++ b/src/halfutils.h @@ -0,0 +1,9 @@ +#ifndef HALFUTILS_H +#define HALFUTILS_H + +#include "halfvec.h" + +double HalfvecL2DistanceSquared(HalfVector * a, HalfVector * b); +double HalfvecInnerProduct(HalfVector * a, HalfVector * b); + +#endif diff --git a/src/halfvec.c b/src/halfvec.c index 3383ff6..7f5687b 100644 --- a/src/halfvec.c +++ b/src/halfvec.c @@ -6,6 +6,7 @@ #include "catalog/pg_type.h" #include "common/shortest_dec.h" #include "fmgr.h" +#include "halfutils.h" #include "halfvec.h" #include "lib/stringinfo.h" #include "libpq/pqformat.h" @@ -800,56 +801,6 @@ vector_to_halfvec(PG_FUNCTION_ARGS) PG_RETURN_POINTER(result); } -/* - * Get the L2 squared distance between half vectors - */ -static double -l2_distance_squared_internal(HalfVector * a, HalfVector * b) -{ - half *ax = a->x; - 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_storeu_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; -} - /* * Get the L2 distance between half vectors */ @@ -862,7 +813,7 @@ halfvec_l2_distance(PG_FUNCTION_ARGS) CheckDims(a, b); - PG_RETURN_FLOAT8(sqrt(l2_distance_squared_internal(a, b))); + PG_RETURN_FLOAT8(sqrt(HalfvecL2DistanceSquared(a, b))); } /* @@ -877,48 +828,7 @@ halfvec_l2_squared_distance(PG_FUNCTION_ARGS) CheckDims(a, b); - PG_RETURN_FLOAT8(l2_distance_squared_internal(a, b)); -} - -/* - * Get the inner product of two half vectors - */ -static double -inner_product_internal(HalfVector * a, HalfVector * b) -{ - half *ax = a->x; - 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); - - dist = _mm256_fmadd_ps(axs, bxs, dist); - } - - _mm256_storeu_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++) - distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]); -#else - /* Auto-vectorized */ - for (int i = 0; i < a->dim; i++) - distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]); -#endif - - return (double) distance; + PG_RETURN_FLOAT8(HalfvecL2DistanceSquared(a, b)); } /* @@ -933,7 +843,7 @@ halfvec_inner_product(PG_FUNCTION_ARGS) CheckDims(a, b); - PG_RETURN_FLOAT8(inner_product_internal(a, b)); + PG_RETURN_FLOAT8(HalfvecInnerProduct(a, b)); } /* @@ -948,7 +858,7 @@ halfvec_negative_inner_product(PG_FUNCTION_ARGS) CheckDims(a, b); - PG_RETURN_FLOAT8(-inner_product_internal(a, b)); + PG_RETURN_FLOAT8(-HalfvecInnerProduct(a, b)); } /*