pgvector 源码学习: 6.4 CPU 优化与调度

pgvector 源码学习: 6.4 CPU 优化与调度

本文介绍 pgvector 如何利用现代 CPU features（CPU 特性）来优化向量距离计算，它通过 runtime CPU feature detection（运行时 CPU 特性检测）和 automatic code dispatching（自动代码调度）实现。系统会在 extension initialization time（扩展初始化时）根据可用的 CPU instructions（CPU 指令）选择最高效的实现。

概览

pgvector 采用两种互补的策略来优化在不同 CPU architectures（CPU 架构）上的性能：

1. Runtime Dispatching（运行时调度）: 对于 half-precision（半精度）向量类型 halfvec 和 bit vector operations（位向量操作），系统会在初始化时检测 CPU features（CPU 特性），并使用 function pointers（函数指针）选择优化的实现。
2. Compile-time Multi-versioning（编译时多版本）: 对于 full-precision（全精度）向量类型 vector 的操作，GCC'starget_clonesattribute（GCC 的 target_clones 属性）会生成函数的多个版本，并由编译器自动调度。

来源:
src/halfutils.c 1-299
src/bitutils.c 1-223
src/vector.c 32-36

调度架构(Dispatching Architecture)

来源:
src/vector.c 48-55
src/halfutils.c 277-298
src/bitutils.c 206-222

运行时调度: 半向量操作(Runtime Dispatching: Half Vector Operations)

CPU 特性检测(CPU Feature Detection)

半向量操作使用 runtime CPU feature detection（运行时 CPU 特性检测）来选择优化的实现。检测会检查对 AVX、F16C 和 FMA 的支持：

来源:
src/halfutils.c 252-298
src/halfutils.h 13-16

函数指针声明(Function Pointer Declarations)

系统使用在初始化期间分配的全局 function pointers（函数指针）：

来源:
src/halfutils.h 13-16
src/halfutils.c 22-25

F16C SIMD 实现(F16C SIMD Implementation)

F16C 优化的实现使用 AVX intrinsics（AVX 内联函数）同时处理 8 个半精度值：

Key operations (关键操作):

• _mm_loadu_si128(): Load 128 bits (8 half values) unaligned (加载 128 位/8 个半精度值，未对齐)
• _mm256_cvtph_ps(): Convert 8 FP16 to 8 FP32 using F16C (使用 F16C 将 8 个 FP16 转换为 8 个 FP32)
• _mm256_fmadd_ps(): Fused multiply-add (融合乘加) on 8 FP32 values (对 8 个 FP32 值执行融合乘加)
• _mm256_storeu_ps(): Store results for horizontal reduction (存储结果以进行水平归约)

Example L2 Distance computation (L2 距离计算示例) (每次迭代处理 8 个维度):

for (i = 0; i < count; i += 8)  
{  
    Load 8 half values from ax[i:i+8] (从 ax[i:i+8] 加载 8 个半精度值)  
    Load 8 half values from bx[i:i+8] (从 bx[i:i+8] 加载 8 个半精度值)  
    Convert both to 8 float32 values (将两者都转换为 8 个 float32 值)  
    diff = ax_float - bx_float  
    distance += diff * diff  (using FMA / 使用 FMA)  
}  
Horizontal sum of 8 accumulated values (对 8 个累加值进行水平求和)  
Process remaining dimensions with scalar code (使用标量代码处理剩余维度)  

来源:
src/halfutils.c 44-76
src/halfutils.c 92-119
src/halfutils.c 145-192

条件编译指令(Conditional Compilation Directives)

F16C 的实现基于以下 Conditional Compilation Directives（条件编译指令）进行条件编译：

来源:
src/halfutils.c 6-20
src/halfutils.c 240-274

运行时调度: 位向量操作(Runtime Dispatching: Bit Vector Operations)

AVX512 检测(AVX512 Detection)

位向量操作检测对 AVX512F 和 AVX512VPOPCNTDQ 的支持：

来源:
src/bitutils.c 171-222

函数指针声明(Function Pointer Declarations)

来源:
src/bitutils.h 11-12
src/bitutils.c 44-45

AVX512 SIMD 实现(AVX512 SIMD Implementation)

AVX512 实现每次迭代处理 64 字节（512 位）：

Key operations (关键操作):

• _mm512_loadu_si512(): Load 512 bits unaligned (加载 512 位，未对齐)
• _mm512_xor_si512(): XOR (异或) two 512-bit vectors (两个 512 位向量)
• _mm512_popcnt_epi64(): Population count (群体计数) on 8×64-bit lanes (对 8×64 位通道执行群体计数)
• _mm512_reduce_add_epi64(): Horizontal sum (水平求和) of 8 64-bit values (8 个 64 位值)

Hamming distance computation (汉明距离计算):

for each 64-byte chunk: (对于每个 64 字节块:)  
    Load 512 bits from ax (从 ax 加载 512 位)  
    Load 512 bits from bx (从 bx 加载 512 位)  
    XOR the vectors (对向量进行异或)  
    POPCNT each 64-bit lane (8 counts) (对每个 64 位通道进行 POPCNT/8 个计数)  
    Accumulate 8 counts (累加 8 个计数)  
Horizontal sum (水平求和)  
Process remaining bytes with default implementation (使用默认实现处理剩余字节)  

来源:
src/bitutils.c 74-93
src/bitutils.c 132-157

POPCNT 的 Target Clones(Target Clones for POPCNT)

默认的位实现使用 GCC'starget_clonesattribute 来进行 POPCNT（群体计数）优化：

#define BIT_TARGET_CLONES __attribute__((target_clones("default", "popcnt")))  

这会生成 BitHammingDistanceDefault 和 BitJaccardDistanceDefault 的两个版本：

• default: Uses lookup table (查找表) pg_number_of_ones[]
• popcnt: Uses CPU POPCNT instruction (CPU POPCNT 指令)

来源:
src/bitutils.c 28-32
src/bitutils.c 47-71

编译时多版本: 向量操作(Compile-time Multi-versioning: Vector Operations)

Target Clones 属性(Target Clones Attribute)

Full-precision vector operations（全精度向量操作）使用 GCC'starget_clones 来生成多个函数版本：

#if defined(USE_TARGET_CLONES) && !defined(__FMA__)  
#define VECTOR_TARGET_CLONES __attribute__((target_clones("default", "fma")))  
#else  
#define VECTOR_TARGET_CLONES  
#endif  

该属性应用于距离函数：

来源:
src/vector.c 32-36
src/vector.c 549-724

FMA 优化(FMA Optimization)

FMA（Fused Multiply-Add，融合乘加）指令将乘法和加法组合在一个操作中，具有更高的精度和 throughput（吞吐量）：

Without FMA (没有 FMA):

for (i = 0; i < dim; i++)  
    distance += ax[i] * bx[i];  // Two operations: mul, add (两个操作: 乘, 加)  

With FMA (有 FMA):

for (i = 0; i < dim; i++)  
    distance = fma(ax[i], bx[i], distance);  // Single FMA operation (单个 FMA 操作)  

编译器会自动 vectorizes（向量化）这些循环，并根据 CPU capabilities（CPU 能力）在 runtime（运行时）选择合适的版本。

来源:
src/vector.c 549-606

自动向量化注释(Auto-vectorization Comments)

代码中包含 /* Auto-vectorized */ 注释，用于指示 GCC/Clang 优化为 SIMD instructions（SIMD 指令）的循环：

VECTOR_TARGET_CLONES static float
VectorL2SquaredDistance(int dim, float *ax, float *bx)  
{  
float distance = 0.0;  

    /* Auto-vectorized */  
for (int i = 0; i < dim; i++)  
    {  
float diff = ax[i] - bx[i];  
        distance += diff * diff;  
    }  

return distance;  
}  

来源:
src/vector.c 549-563
src/halfvec.c 696-702

半精度转换优化(Half-Precision Conversion Optimization)

内联转换函数(Inline Conversion Functions)

HalfToFloat4() 和 Float4ToHalf() 函数是关键的 hot paths（热路径），具有多种实现策略：

来源:
src/halfutils.h 62-141
src/halfutils.h 146-261

转换实现比较(Conversion Implementation Comparison)

软件实现处理以下特殊情况：

• Sign bit extraction and positioning (符号位提取和定位)
• Exponent bias conversion (FP16: -15, FP32: -127) (指数偏差转换)
• Mantissa scaling (尾数缩放)
• Special cases: infinity (无穷大), NaN (非数字), subnormal numbers (非规范化数)
• Rounding (舍入) for FP32→FP16 conversion

来源:
src/halfutils.h 59-263

构建时特性检测(Build-time Feature Detection)

条件编译宏(Conditional Compilation Macros)

Build system（构建系统）定义 macros（宏）来启用/禁用优化：

来源:
src/halfutils.c 6-20
src/bitutils.c 7-9
src/vector.c 32-36

平台特定编译(Platform-specific Compilation)

CPUID Access (CPUID 访问):

• GCC/Clang: __get_cpuid() from <cpuid.h> when USE__GET_CPUID defined
• MSVC: __cpuid() from <intrin.h>

Target Attributes (目标属性):

• GCC/Clang: __attribute__((target(...)))
• MSVC: Empty macro (空宏) (MSVC allows intrinsics (内联函数) without attributes)

来源:
src/halfutils.c 9-13
src/bitutils.c 14-18

性能特征(Performance Characteristics)

优化带来的加速(Speedup from Optimizations)

这些优化提供了显著的性能提升：

来源:
src/halfutils.c 27-238
src/bitutils.c 47-157

吞吐量分析(Throughput Analysis)

Half Vector Operations (F16C) (半向量操作 - F16C):

• Loads 8 FP16 values (128 bits) (加载 8 个 FP16 值/128 位)
• Converts to 8 FP32 values in parallel (并行转换为 8 个 FP32 值)
• Processes 8 dimensions per inner loop iteration (每次内循环迭代处理 8 个维度)
• Remainder (剩余部分) handled with scalar code (标量代码处理)

Bit Vector Operations (AVX512) (位向量操作 - AVX512):

• Processes 512 bits (64 bytes) per iteration (每次迭代处理 512 位/64 字节)
• 8×64-bit POPCNT operations in parallel (并行执行 8×64 位 POPCNT 操作)
• Remainder (剩余部分) handled by default implementation (默认实现处理)

来源:
src/halfutils.c 44-76
src/bitutils.c 74-93

扩展初始化流程(Extension Initialization Flow)

来源:
src/vector.c 48-55
src/halfutils.c 277-298
src/bitutils.c 206-222

代码组织(Code Organization)

关键源文件(Key Source Files)

来源:
src/vector.c 1-1322
src/halfutils.c 1-299
src/halfutils.h 1-264
src/bitutils.c 1-223
src/bitutils.h 1-17

函数命名约定(Function Naming Conventions)

Runtime Dispatched Functions (运行时调度的函数):

• {Type}{Operation}Default: Portable C implementation (可移植 C 实现)
• {Type}{Operation}F16c: F16C/AVX optimized implementation (F16C/AVX 优化实现)
• {Type}{Operation}Avx512Popcount: AVX512 optimized implementation (AVX512 优化实现)

Examples (示例):

• HalfvecL2SquaredDistanceDefault vs HalfvecL2SquaredDistanceF16c
• BitHammingDistanceDefault vs BitHammingDistanceAvx512Popcount

Target Clones Functions (Target Clones 函数):

• Standard function name (标准函数名) (e.g., VectorL2SquaredDistance)
• Compiler (编译器) generates {name}.default and {name}.fma variants (生成变体)
• IFUNC resolver (IFUNC 解析器) selects at runtime (运行时)

来源:
src/halfutils.c 27-238
src/bitutils.c 47-157
src/vector.c 549-724