Add AVX-512 FP16 implementation of halfvec distance functions #531

lucagiac81 · 2024-04-24T14:28:09Z

This PR adds implementations of halfvec distance functions based on the AVX-512 FP16 instruction set. The instruction set was introduced with Intel 4th Gen Intel® Xeon® Scalable processors. It supports 32x FP16 operations per instruction with 512-bit registers.

Compiler support for the new instructions was added in gcc-12 and clang-14. Those versions are minimum requirements for the AVX-512 FP16 functions to be compiled (controlled by conditional compilation). Support for the instruction set is also detected at runtime using CPUID. If not supported, the existing default or F16c functions are used.

Building was tested with

gcc-11/clang-13 (no AVX-512 FP16 support)
gcc-12/clang-14 (with AVX-512 FP16 support)

Execution of a binary compiled with gcc-12 (which includes the AVX-512 FP16 functions) was tested on

4th Gen Intel® Xeon® Scalable processor (with AVX-512 FP16 support): AVX-512 FP16 functions are used
3rd Gen Intel® Xeon® Scalable processor (no AVX-512 FP16 support): existing F16c functions are used

There is one open question regarding CI with GitHub Actions. The AVX-512 FP16 functions will not be tested unless a runner supports the instruction set.

Performance results will be shared soon.

jkatz · 2024-04-24T14:45:19Z

@nathan-bossart Would love your feedback on this one.

nathan-bossart

Performance results will be shared soon.

Looking forward to these!

nathan-bossart · 2024-04-24T16:28:26Z

src/halfutils.c

+#ifdef HAVE_AVX512FP16
+TARGET_XSAVE static bool
+SupportsAvx512Fp16()
+{
+ unsigned int exx[4] = {0, 0, 0, 0};
+ unsigned int feature = (1 << 23);
+
+#if defined(HAVE__GET_CPUID)
+ __get_cpuid_count(7, 0, &exx[0], &exx[1], &exx[2], &exx[3]);
+#elif defined(HAVE__CPUID)
+ __cpuid(exx, 7, 0);
+#endif
+
+ return (exx[3] & feature) == feature;
+}
+#endif


I think this is missing a couple steps, such as checking for osxsave and verifying the ZMM registers are enabled. See SupportsAvx512Popcount() for an example.

Thanks for the reference. I'll add those checks (OSXSAVE and XCR0 control register).

nathan-bossart · 2024-04-24T16:29:28Z

src/halfutils.c

+ for (; i < dim; i++)
+ distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);


Is this auto-vectorized? (Same question for HalfvecL2SquaredDistanceAvx512Fp16().)

I checked L2SquaredDistance and InnerProduct, and it is using AVX scalar instructions, at least with gcc-12. We'll try masked vector instructions to handle the loop remainder.

The latest update includes masked vector instructions for the loop remainder.

nathan-bossart · 2024-04-24T16:32:39Z

src/halfutils.c

+#ifdef HAVE_AVX512FP16
+ if (SupportsAvx512Fp16())
+ {
+ HalfvecL2SquaredDistance = HalfvecL2SquaredDistanceAvx512Fp16;
+ HalfvecInnerProduct = HalfvecInnerProductAvx512Fp16;
+ HalfvecCosineSimilarity = HalfvecCosineSimilarityAvx512Fp16;
+ }
+#endif


nitpick: This might not need to be nested in the HALFVEC_DISPATCH block.

You're right. Currently, it is taking advantage of the OSXSAVE check included with the other features, but I'll separate that.

jkatz · 2024-04-24T16:49:06Z

I'll kick off some local benchmark runs to see the diffs. I have a r7i at the ready.

jkatz · 2024-04-24T17:34:11Z

@lucagiac81 I'm having issues compiling on an EC2 r7i. This is using gcc12 and clang-15. Here is some truncated output:

/usr/bin/clang-15 -Wno-ignored-attributes -fno-strict-aliasing -fwrapv -Xclang -no-opaque-pointers -Wno-unused-command-line-argument -Wno-compound-token-split-by-macro -O2  -I. -I./ -I/usr/include/postgresql/16/server -I/usr/include/postgresql/internal  -Wdate-time -D_FORTIFY_SOURCE=2 -D_GNU_SOURCE -I/usr/include/libxml2  -flto=thin -emit-llvm -c -o src/halfutils.bc src/halfutils.c
src/halfutils.c:92:9: error: expected ';' after expression
        __m512h         dist = _mm512_setzero_ph();
               ^
               ;
src/halfutils.c:92:2: error: use of undeclared identifier '__m512h'
        __m512h         dist = _mm512_setzero_ph();
        ^
src/halfutils.c:92:11: error: use of undeclared identifier 'dist'
        __m512h         dist = _mm512_setzero_ph();
                        ^
src/halfutils.c:92:18: warning: call to undeclared function '_mm512_setzero_ph'; ISO C99 and later do not support implicit function declarations [-Wimplicit-function-declaration]
        __m512h         dist = _mm512_setzero_ph();
                               ^
src/halfutils.c:95:10: error: expected ';' after expression
                __m512h axi = _mm512_loadu_ph(ax+i);
                       ^
                       ;
src/halfutils.c:95:3: error: use of undeclared identifier '__m512h'
                __m512h axi = _mm512_loadu_ph(ax+i);
                ^
src/halfutils.c:95:11: error: use of undeclared identifier 'axi'; did you mean 'ax'?
                __m512h axi = _mm512_loadu_ph(ax+i);
                        ^~~
                        ax
src/halfutils.c:87:52: note: 'ax' declared here
HalfvecL2SquaredDistanceAvx512Fp16(int dim, half * ax, half * bx)

lucagiac81 · 2024-04-25T02:49:52Z

@jkatz I think clang is not applying __attribute__((target("avx512fp16")))

I tested on an m7i instance (where -march=native includes -mavx512fp16) with clang-15

Build pgvector with the default Makefile (which has -march=native in OPTFLAGS): no error
Remove -march=native from Makefile: compilation errors (similar to your report)
Replace -march=native with -mavx512fp16 in Makefile: no error

With gcc-12.3, and I got no errors in all cases.

Can you try adding -march=native or -mavx512fp16 to your flags as a temporary solution?

lucagiac81 · 2024-04-25T19:54:40Z

Rebased on latest master
Added checks for OSXSAVE and zmm registers enabled
Added L1 distance AVX512-FP16 implementation

akashsha1 · 2024-04-26T23:52:17Z

src/halfutils.c

+SupportsAvx512Fp16()
+{
+ unsigned int exx[4] = {0, 0, 0, 0};
+ unsigned int feature = (1 << 23);


nit. feature can be defined using #DEFINE CPU_FEATURE_AVX512FP16

akashsha1 · 2024-04-26T23:54:23Z

src/halfutils.c

+ __cpuid(exx, 7, 0);
+#endif
+
+ /* Check OS supports XSAVE */


nit. update comment to reflect OSXSAVE

akashsha1 · 2024-04-27T00:07:48Z

src/halfutils.c

+ return false;
+
+ /* Check XMM, YMM, and ZMM registers are enabled */
+ if ((_xgetbv(0) & 0xe6) != 0xe6)


@nathan-bossart shouldn't this be _xgetbv(0) & 0xe6) == 0xe6 ? Similar comment on L187 in bitutils.c per the discussion [0]

[0] : https://www.postgresql.org/message-id/20240418210158.GA3776258%40nathanxps13

This looks alright as-is to me. If this check fails, we return false, so != looks correct.

lucagiac81 · 2024-05-14T03:18:13Z

While collecting data with ANN benchmarks, we noticed a degradation in recall for some datasets (such as sift-128) when computing distances in half precision. Other datasets (such as gist-960) are not affected, and recall is matched to the existing distance functions. The existing functions (*F16c) first convert halfvec elements to single precision and execute the distance computation in single precision.

So, enabling the FP16 distance functions may not be desirable in all cases. The latest update to the PR provides two implementations of the distance functions with AVX-512: one using single precision and one using half precision.

The single-precision functions are used by default, as they're the most generically applicable.
The user can decide to use the half-precision functions by setting a variable (halfvec.use_fp16_compute). The decision should be based on recall vs performance tradeoff for a specific dataset (this implementation handles 2x the vector elements per iteration compared to the single-precision one).

jkatz · 2024-05-14T15:30:18Z

@lucagiac81 Thanks for the continued work. Per @nathan-bossart comment earlier, it'd be helpful to see the actual performance results.

I'll try to get this to build again - last I checked I didn't have avx512fp16 available on my instance class.

lucagiac81 · 2024-05-16T19:36:13Z

Here are some initial results

We made a few changes to ANN benchmarks to support halfvec (following your post) and used the default conditions.
The tests were run on an m7i.metal-24xl instance. pgvector is compiled with gcc-12.
We measure query performance for the AVX-512 half-precision and single-precision distance implementations, and compare with the existing F16c implementation. To enable the half-precision functions, we use SET halfvec.use_fp16_compute = true in set_query_arguments, as described in the previous comment.

With the gist-960-euclidean dataset, so far we observe

qps increase of 9.6%-12.9% for half-precision and 1.4%-3.6% for single-precision
p99 reduction of 8.3%-12.8% for half-precision and 0-4% for single-precision
recall is matched to within +/-1% for both implementations

It'd be great if you could reproduce these numbers with your setup. Please let me know if you still run into compilation issues. We'd also like to collect data with dbpedia-openai-1000k-angular as well (higher dimensions, different distance metric) , but we're running into a 403 error when downloading the dataset (similar to this report). Do you have any advice on how to run with that dataset?

nathan-bossart reviewed Apr 24, 2024

View reviewed changes

lucagiac81 force-pushed the halfvec_avx512_fp16 branch from e390649 to 85ba2dc Compare April 25, 2024 19:53

lucagiac81 marked this pull request as ready for review April 25, 2024 19:54

akashsha1 reviewed Apr 26, 2024

View reviewed changes

src/halfutils.c Outdated

__cpuid(exx, 7, 0);

#endif

/* Check OS supports XSAVE */

Copy link

akashsha1 Apr 26, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit. update comment to reflect OSXSAVE

akashsha1 reviewed Apr 27, 2024

View reviewed changes

lucagiac81 added 4 commits May 13, 2024 16:22

Implement AVX512 distance calculations for halfvec

061d1ce

Use masked instructions for remainder loops

e43ae81

Add option to use fp16 computation

fe7876c

Add tests for fp16 distance functions

915d6eb

lucagiac81 force-pushed the halfvec_avx512_fp16 branch from 85ba2dc to 915d6eb Compare May 14, 2024 03:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add AVX-512 FP16 implementation of halfvec distance functions #531

Add AVX-512 FP16 implementation of halfvec distance functions #531

lucagiac81 commented Apr 24, 2024

jkatz commented Apr 24, 2024

nathan-bossart left a comment

nathan-bossart Apr 24, 2024

lucagiac81 Apr 25, 2024

nathan-bossart Apr 24, 2024

lucagiac81 Apr 25, 2024

lucagiac81 May 14, 2024

nathan-bossart Apr 24, 2024

lucagiac81 Apr 25, 2024

jkatz commented Apr 24, 2024

jkatz commented Apr 24, 2024 •

edited

lucagiac81 commented Apr 25, 2024 •

edited

lucagiac81 commented Apr 25, 2024

akashsha1 Apr 26, 2024 •

edited

akashsha1 Apr 26, 2024

akashsha1 Apr 27, 2024 •

edited

nathan-bossart Apr 29, 2024

lucagiac81 commented May 14, 2024

jkatz commented May 14, 2024

lucagiac81 commented May 16, 2024

		for (; i < dim; i++)
		distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);

Add AVX-512 FP16 implementation of halfvec distance functions #531

Are you sure you want to change the base?

Add AVX-512 FP16 implementation of halfvec distance functions #531

Conversation

lucagiac81 commented Apr 24, 2024

jkatz commented Apr 24, 2024

nathan-bossart left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkatz commented Apr 24, 2024

jkatz commented Apr 24, 2024 • edited

lucagiac81 commented Apr 25, 2024 • edited

lucagiac81 commented Apr 25, 2024

akashsha1 Apr 26, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

akashsha1 Apr 27, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lucagiac81 commented May 14, 2024

jkatz commented May 14, 2024

lucagiac81 commented May 16, 2024

jkatz commented Apr 24, 2024 •

edited

lucagiac81 commented Apr 25, 2024 •

edited

akashsha1 Apr 26, 2024 •

edited

akashsha1 Apr 27, 2024 •

edited