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

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

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

@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)

@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?

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

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).

@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.

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?

Hi @lucagiac81, thanks for the PR, and sorry for the delay. Based on the numbers above, I'm not sure the benefit justifies the complexity.

For the dbpedia-openai-1000k-angular dataset, there's a create_dataset.py script in ann-benchmarks, fwiw.

Thanks @ankane. The issue with the dbpedia dataset was the inability to create the dataset locally. After updating the datasets package to v2.19.1, the issue is resolved. We'll share the results with that dataset as well.

Regarding the complexity, is the additional parameter to select the precision of the distance computation the main concern?

Sounds good.

The parameter doubles my concern, but there's still a lot of complexity without it. I'm not sure either choice is great based on the numbers above, since the single-precision version provides little performance benefit and while the half-precision provides some benefit but reduces precision and range.

For comparison, here are the f16c + fma numbers: #311 (comment).

The latest update eliminates the need to manually enable FP16 computation. Computation starts with FP16 and switches to FP32 in case of overflow.
With this approach, the recall degradation we observed with FP16 for certain datasets (sift-128-euclidean) is resolved. Datasets that performed well with FP16 computation (gist-960-euclidean) show similar performance as with FP16 enabled in the previous version of the code.
We will share a more detailed report, but this is a step towards reducing complexity (especially not exposing it to the end user).

The latest update introduces an AVX512_FP16 implementation of vector_to_halfvec conversion (in separate commit), as that function has a noticeable contribution in VectorDBBench benchmarks. It is also rebased on pgvector v0.8.0. Sharing additional performance measurements with ANN-Benchmarks and VectorDBBench below.

For ANN-Benchmarks, we used a similar setup as for the previously shared data. With the latest changes, the manual selection of FP16/FP32 computation is removed. We also include measurements for index build time (with 8 parallel workers). Performance gains are relative to the existing F16c implementation.
For dbpedia-openai-1000k-angular dataset, we observe

• qps gain of 9.4%-14.2%
• p99 reduction of 9.0%-12.9%
• index build time reduction of 3.8% (m=16) and 11.6% (m=24)
• recall matched within +/-0.3%

For sift-128-euclidean, we observe 11-12% index build time reduction with recall matched within +/-0.1%, but no significant qps/p99 gain. This confirms that the recall degradation previously observed with FP16 computation is resolved and there are gains for certain metrics even at lower dimensions.

For VectorDBBench, we focused on larger datasets and varying search concurrency (1-40 range on an r7i.12xlarge instance). Below are initial results with two datasets:

• Performance1536D5M (OpenAI, 5M vectors, 1536 dimensions): 4-8% qps gain
• Performance768D10M (Cohere, 10M vectors, 768 dimensions): 5-6% qps gain

The last update fixes some issues reported by CI:

• Remove unnecessary variable initialization
• Add AVX512DQ check (required by one instruction)
• Increase clang minimum version requirement from 14 to 16. Version 16 solves the issue reported in a previous comment. I'll keep checking for a solution with earlier versions.

@lucagiac81

I tested the latest revision a1e3ead using vectordbbench on a r7i.8xlarge instance and achieved qps improvements ranging from 7% to 23%, depending on the dataset.

Tests were performed using:

• ef_construction 128
• ef_search 128
• k 100
• full vector stored in the table
• Amazon linux2
• gcc 12.4.0
• binutils 2.39

Note: gcc >= 12 & binutils >= 2.38 are required for this change to have any effect. Upgrading gcc, without binutils can cause pgvector's make to fail.

@lucagiac81

I tested the latest revision a1e3ead using vectordbbench on a r7i.8xlarge instance and achieved qps improvements ranging from 7% to 23%, depending on the dataset.

Dataset qps recall with avx512 qps qps improvement with avx512 recall recallimpact
Performance1536D50K 10644.763 0.9691 12594.4903 18.32% 0.9626 -0.67%
Performance1536D5M 6810.1588 0.9338 7332.7366 7.67% 0.927 -0.73%
Performance768D10M 7916.9702 0.8975 9799.3695 23.78% 0.8949 -0.29%
Performance768D1M 9847.2079 0.926 11043.1437 12.14% 0.9262 0.02%
Tests were performed using:

• ef_construction 128
• ef_search 128
• k 100
• full vector stored in the table
• Amazon linux2
• gcc 12.4.0
• binutils 2.39

Note: gcc >= 12 & binutils >= 2.38 are required for this change to have any effect. Upgrading gcc, without binutils can cause make to fail.

@greenhal thanks for sharing these. This is great improvements.
@jkatz , @nathan-bossart , @ankane - 23% improvements seem like a solid improvement on AWS instances to consider. The PR has also been simplified to automatically take advantage of the new instructions in avx512. Would you be able to help provide any feedback here, and help merge this change?

We'd like to propose a refactoring to make platform-specific optimizations easier to maintain.

In the latest update, AVX-512 functions are moved to separate files. The functions are also included by conditional compilation (enabled by default). In this way, AVX-512 implementations are all in one place, and they can be easily disabled in the build if desired.
AVX-512 feature detection is currently in the same files. In the future, we could also move detection to a common location if duplication becomes a concern.

This approach is easily extended to future contributions and different architectures. This will allow pgvector to benefit from targeted optimizations while keeping the "core" code uncluttered. It will also clarify where particular expertise is needed for maintenance and improvements.

@ankane do you see this as a viable approach to integrate these optimizations? Please let us know your thoughts.

Hi @lucagiac81, thanks again for the PR, but I don't think it's a good fit for pgvector. Changing the accumulation precision can reduce recall (as @greenhal's tests show), and there is already a separate lever for trading recall for speed (ef_search).