Мой SIMD быстрее вашего

Нерассказанная история о том, как мы делаем векторный поиск в LanceDB быстрым.

В январе команда инженеров Eto приняла решение переписать колоночный формат Lance на Rust. Хотя это решение может показаться обычной тенденцией Rewrite-X-In-Rust, за ним стоит нерассказанная история: нашей целью было добиться отличной производительности в недавно анонсировали SSD-устойчивый векторный поиск в LanceDB. Родная генерация кода Rust, отличный набор инструментов для настройки производительности ( cargo asm / cargo flamegraph ), а также огромное количество первоклассных встроенных функций ЦП/поддержка SIMD в стандартной библиотеке (std::arch ) сильно повлияли на наше решение.

Самый фундаментальный запрос в векторном поиске — это «поиск K ближайших соседей (KNN) в многомерном векторном пространстве». Для выполнения этого поиска необходимо вычислить расстояния между векторами. В этой статье мы будем использовать в качестве примера классическое евклидово расстояние (L2). Другие популярные меры расстояния включают косинусное расстояние, скалярное произведение и расстояние Хэмминга.

Первая попытка: наивная реализация и пусть компилятор сделает это

Расстояние L2 довольно просто реализовать в Rust.

pub fn l2(x: &[f32], y: &[f32]) -> f32 {
    x.iter()
        .zip(y.iter())
        .map(|(a, b)| (a - b).powi(2))
        .sum::<f32>()
}

Вычисление L2 между одним вектором запроса (размер 1024) и 1 миллионом векторов занимает 1.25s в экземпляре, оптимизированном для вычислений AWS (C6in.4xlarge).

Компилятор Rust rustc (1.68) выполняет довольно приличную генерацию кода для приведенного выше кода, как показано в Compiler Explorer с RUSTFLAGS=-C opt-level=3 -C target-cpu=icelake-server -C target-feature=+avx2: запускается развертывание цикла:

/// RUSTFLAGS="-C opt-level=3 -C target-cpu=icelake-server"
...
.LBB0_5:
        vmovss  xmm1, dword ptr [rdi + 4*rsi]
        vmovss  xmm2, dword ptr [rdi + 4*rsi + 4]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vsubss  xmm1, xmm2, dword ptr [rdx + 4*rsi + 4]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 8]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 8]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 12]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 12]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 16]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 16]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 20]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 20]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 24]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 24]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        vmovss  xmm1, dword ptr [rdi + 4*rsi + 28]
        lea     rax, [rsi + 8]
        vsubss  xmm1, xmm1, dword ptr [rdx + 4*rsi + 28]
        vmulss  xmm1, xmm1, xmm1
        vaddss  xmm0, xmm0, xmm1
        mov     rsi, rax
        cmp     rcx, rax
        jne     .LBB0_5
...

Люди часто говорят: Пусть компилятор оптимизирует код за вас. К сожалению, vsubss/vaddss/vmulss — это скалярные инструкции для X86_64. Похоже, что LLVM не выполняет автоматическую векторизацию цикла. Было обнаружено, что для LLVM/GCC было бы довольно сложно векторизовать циклы без статической длины во время компиляции. Но наш векторный индекс ДОЛЖЕН поддерживать любое измерение пользовательских векторов, неужели мы оставили какую-то производительность на столе?

Вторая попытка: вычислительное ядро ​​Arrow

LanceDB построен на основе Apache Arrow (Rust). Можем ли мы добиться большего успеха, напрямую используя вычислительные ядра Arrow?

// RUSTFLAGS="-C target-cpu=native -C target-feature=+avx2"
// Use arrow arith kernels.
use arrow_arith::arithmetic::{multiply, subtract};
use arrow_arith::arity::binary;

#[inline]
fn l2_arrow_1(x: &Float32Array, y: &Float32Array) -> f32 {
    let s = subtract(x, y).unwrap();
    let m = multiply(&s, &s).unwrap();
    sum(&m).unwrap()
}

#[inline]
fn l2_arrow_2(x: &Float32Array, y: &Float32Array) -> f32 {
    let m: Float32Array = binary(x, y, |a, b| (a - b).powi(2)).unwrap();
    sum(&m).unwrap()
}

Работая на той же машине c6in.4xlarge с RUSTFLAGS=”-C target-cpu=native -C target-feature=+avx2” , мы получили 2.81s и 2.16s соответственно. Удивительно, но использование ядра вычислений со стрелками работает медленнее, чем наша наивная реализация. Причина в том, что, в отличие от встроенной функции L2, которая имеет нулевое выделение памяти и сканирует только x и y один раз, что делает недействительным кэш L1 только тогда, когда это абсолютно необходимо. l2_arrow_1 требует однократного сканирования массива x, y, m, s, дважды аннулирования кеша L1/L2 и двух дополнительных выделений памяти для s и m. l2_arrow_2 сканирует x, yonce, что улучшает поведение кэша и требует только одного дополнительного выделения памяти.

Третья попытка: BLAS спасти?

Многие библиотеки научных вычислений используют BLAS для ускорения матричных/векторных вычислений. LanceDB в некоторых местах также использует BLAS!

В частности, numpy считается хорошо оптимизированным для алгоритмов линейной алгебры. Мы используем Numpy + Intel MKL, одну из самых быстрых библиотек BLAS для оценки производительности BLAS.

>>> import numpy as np
>>> np.show_config()
blas_armpl_info:
  NOT AVAILABLE
blas_mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/include']
blas_opt_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/include']
lapack_armpl_info:
  NOT AVAILABLE
lapack_mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/include']
lapack_opt_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/home/ubuntu/miniconda3/envs/np-mkl/include']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2,AVX512F,AVX512CD,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL
    not found = AVX512_KNL,AVX512_KNM

Запуск простого скрипта Python на том же экземпляре AWS c6in занимает 5.498s

import numpy as np
import time

x = np.random.random((1, 1024))
y = np.random.random((1024*1024, 1024))

total = 10
start = time.time()
for i in range(total):
    d = np.linalg.norm(y - x, axis=1)
print("time: {}s".format((time.time() - start) / total))

Последняя попытка: можем ли мы выжать из современного процессора каждую копейку?

Наша философия построения бессерверной базы данных Vector DB LanceDB заключается в полном использовании возможностей современного оборудования.

Современные процессоры необычайно эффективны. Мы твердо верим, что можем лучше использовать ЦП.

Мы переписываем вычисление расстояния L2, используя инструкции Intel AVX-2 в std::arch::x86_64.

use std::arch::x86_64::*;

#[inline]
pub(crate) fn l2_f32(from: &[f32], to: &[f32]) -> f32 {
    unsafe {
        // Get the potion of the vector that is aligned to 32 bytes.
        let len = from.len() / 8 * 8;
        let mut sums = _mm256_setzero_ps();
        for i in (0..len).step_by(8) {
            let left = _mm256_loadu_ps(from.as_ptr().add(i));
            let right = _mm256_loadu_ps(to.as_ptr().add(i));
            let sub = _mm256_sub_ps(left, right);
            // sum = sub * sub + sum
            sums = _mm256_fmadd_ps(sub, sub, sums);
        }
        // Shift and add vector, until only 1 value left.
        // sums = [x0-x7], shift = [x4-x7]
        let mut shift = _mm256_permute2f128_ps(sums, sums, 1);
        // [x0+x4, x1+x5, ..]
        sums = _mm256_add_ps(sums, shift);
        shift = _mm256_permute_ps(sums, 14);
        sums = _mm256_add_ps(sums, shift);
        sums = _mm256_hadd_ps(sums, sums);
        let mut results: [f32; 8] = [0f32; 8];
        _mm256_storeu_ps(results.as_mut_ptr(), sums);

        // Remaining unaligned values
        results[0] += l2_scalar(&from[len..], &to[len..]);
        results[0]
    }
}

На этот раз требуется всего 0.325s для вычисления расстояний L2 по 1 миллиону векторов, что на 350% быстрее, чем у ближайших альтернатив!

Мы не забываем тебя, Apple Silicon!

Наша команда разделяет любовь к MacBook M1/M2 Apple Silicon. Мы стремимся сделать LanceDB лучшим в своем классе вариантом производительности на машинах, которые мы используем каждый день.

Чипы Apple M1/M2 основаны на архитектуре Arm aarch64 с NEON инструкциями. К сожалению, расширение Arm Scalable Vector Extension (SVE) для этих чипов пока недоступно. Итак, мы реализовали только NEON версию L2 distance.

use std::arch::aarch64::*;

#[inline]
pub(crate) fn l2_f32(from: &[f32], to: &[f32]) -> f32 {
    unsafe {
        let len = from.len() / 4 * 4;
        let buf = [0.0_f32; 4];
        let mut sum = vld1q_f32(buf.as_ptr());
        for i in (0..len).step_by(4) {
            let left = vld1q_f32(from.as_ptr().add(i));
            let right = vld1q_f32(to.as_ptr().add(i));
            let sub = vsubq_f32(left, right);
            sum = vfmaq_f32(sum, sub, sub);
        }
        let mut sum = vaddvq_f32(sum);
        sum += l2_scalar(&from[len..], &to[len..]);
        sum
    }
}

Эта реализация L2 с ускорением NEON требует всего 0.299s для вычисления 1 миллиона расстояний на M2 Max MacBook Pro.

Для полноты бенчмаркинга мы также запустили тот же скрипт numpy на Macbook Pro, используя Apple Accelerate Framework для BLAS. Мы видели аналогичное ускорение в 3–15 раз при ручной настройке инструкций SIMD для разных архитектур ЦП ( x86_64 и apple-silicon (aarch64) ).

Как будет выглядеть будущее LanceDB

Несмотря на то, что мы уже значительно ускорили векторные вычисления с помощью нашей реализации SIMD, мы можем сделать гораздо больше, чтобы LanceDB работала еще быстрее:

• Добавьте подпрограммы AVX-512 для расширения векторизации битовой ширины.
• Поддержка векторов bf16, f64floats.
• CUDA и GPU-ускорение Apple Silicon.
• Повышение эффективности кэширования кода PQ/OPQ.

Как только проект std::simd/portable SIMD станет доступен в стабильной версии Rust, мы обязательно перейдем на использование std::simd для простоты.

Кроме того, поскольку LanceDB является устойчивой к SSD и бессерверной векторной базой данных, существует множество дополнительных оптимизаций путей ввода-вывода, которые можно выполнить специально для твердотельных накопителей NVME.

Богатая экосистема Rust дает нам огромные возможности для написания кросс-платформенного аппаратного ускорения с легкостью, за что мы особенно благодарны сообществу Rust за:

• std::arch::{x86_64, aarch64} охватывает большинство наших вариантов использования стабильной цепочки инструментов Rust.
• cargo bench + pprof (https://github.com/tikv/pprof-rs) упрощают поддержку и проверку тестов. Это основа нашей непрерывной работы по сравнительному анализу.
• cargo asm(https://github.com/pacak/cargo-show-asm) для проверки генерации кода на процессорах x86_64 и aarch64, чтобы мы могли проверить, активны ли автовекторизация и другие оптимизации компилятора.
• Compiler Explorer (https://godbolt.org/) — еще один инструмент, позволяющий нам лучше понять оптимизацию компилятора.
• cargo flamegraph (https://github.com/flamegraph-rs/flamegraph) упрощает работу с flamegraph. Если вы уже использовали flamegraph.pl в проекте C/C++, я уверен, что вы искренне оцените проект flamegraph-rs.

Попробуйте LanceDB!

Если вам нравится идея LanceDB , вы можете попробовать ее сегодня, просто pip install lancedb . И, пожалуйста, поставьте нам 🌟 на нашем Github Repo!

Хотите узнать больше о Lance и LanceDB? Пожалуйста, ознакомьтесь с другими интересными статьями: