| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * simd_utils.h - Shared SIMD utilities and patterns | ||
| 6 | * | ||
| 7 | * Provides reusable SIMD infrastructure for distance computations, | ||
| 8 | * quantization, and other vectorized operations. This header contains: | ||
| 9 | * | ||
| 10 | * - Horizontal reduction functions for all SIMD variants | ||
| 11 | * - Dispatch macro templates for runtime CPU detection | ||
| 12 | * - Common SIMD patterns (prefetch, alignment) | ||
| 13 | * - Generic batch processing helpers | ||
| 14 | * | ||
| 15 | * Design rationale: By centralizing SIMD utilities, we enable ~50-60% code | ||
| 16 | * reuse across full-precision, binary, and product quantization distance | ||
| 17 | * implementations. Each distance type needs different intrinsics, but shares | ||
| 18 | * the infrastructure. | ||
| 19 | */ | ||
| 20 | |||
| 21 | #ifndef VS_SIMD_UTILS_H | ||
| 22 | #define VS_SIMD_UTILS_H | ||
| 23 | |||
| 24 | /* Must be first - defines VS_SIMD_NONE used by VS_TARGET_CLONES */ | ||
| 25 | #include "vs_config.h" | ||
| 26 | |||
| 27 | #include <stdatomic.h> | ||
| 28 | #include <stdint.h> | ||
| 29 | |||
| 30 | #include "core/platform.h" | ||
| 31 | |||
| 32 | /* | ||
| 33 | * Target Attribute Macros | ||
| 34 | * | ||
| 35 | * These macros specify the CPU features required for each SIMD implementation. | ||
| 36 | * Using macros makes it easy to change target features in one place. | ||
| 37 | */ | ||
| 38 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 39 | #define VS_TARGET_AVX512 __attribute__((target("avx512f,avx512dq"))) | ||
| 40 | #define VS_TARGET_AVX512_VPOPCNTDQ \ | ||
| 41 | __attribute__((target("avx512f,avx512vpopcntdq"))) | ||
| 42 | #define VS_TARGET_AVX2 __attribute__((target("avx2,fma"))) | ||
| 43 | #define VS_TARGET_F16C_AVX2 __attribute__((target("avx2,fma,f16c"))) | ||
| 44 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 45 | /* NEON is always available on AArch64, no attribute needed */ | ||
| 46 | #define VS_TARGET_NEON | ||
| 47 | #endif | ||
| 48 | |||
| 49 | /* | ||
| 50 | * Target Clones Macro | ||
| 51 | * | ||
| 52 | * Generates multiple function versions for different ISAs. The dynamic linker | ||
| 53 | * selects the best version at load time. Only effective on x86 with GCC 6+ or | ||
| 54 | * Clang 13+. | ||
| 55 | * | ||
| 56 | * Note: target_clones doesn't work effectively on ARM (generates single | ||
| 57 | * clone). | ||
| 58 | */ | ||
| 59 | #ifndef __has_attribute | ||
| 60 | #define __has_attribute(x) 0 | ||
| 61 | #endif | ||
| 62 | |||
| 63 | /* | ||
| 64 | * Disable target_clones when: | ||
| 65 | * - simd=none (VS_SIMD_NONE) | ||
| 66 | * - coverage build (VS_COVERAGE) - each clone is separate, only one executes | ||
| 67 | * - compiler lacks target_clones support | ||
| 68 | * - non-x86 architecture (ARM target_clones generates single clone anyway) | ||
| 69 | */ | ||
| 70 | #if !defined(VS_SIMD_NONE) && !defined(VS_COVERAGE) && \ | ||
| 71 | __has_attribute(target_clones) && defined(__x86_64__) | ||
| 72 | /* | ||
| 73 | * v3 is 256-bit AVX2+FMA, v4 is 512-bit AVX-512+FMA. GCC 11's | ||
| 74 | * target_clones dispatcher rejects those names. haswell and | ||
| 75 | * skylake-avx512 compile there, but GCC tunes the latter to 256-bit | ||
| 76 | * and vectorizes the former less, so they are only the GCC 11 fallback. | ||
| 77 | */ | ||
| 78 | #if defined(__clang__) || __GNUC__ >= 12 | ||
| 79 | #define VS_TARGET_CLONES \ | ||
| 80 | __attribute__(( \ | ||
| 81 | target_clones("default", "arch=x86-64-v3", "arch=x86-64-v4"))) | ||
| 82 | #else | ||
| 83 | #define VS_TARGET_CLONES \ | ||
| 84 | __attribute__(( \ | ||
| 85 | target_clones("default", "arch=haswell", "arch=skylake-avx512"))) | ||
| 86 | #endif | ||
| 87 | #else | ||
| 88 | #define VS_TARGET_CLONES | ||
| 89 | #endif | ||
| 90 | |||
| 91 | /* | ||
| 92 | * Horizontal Reduction Functions | ||
| 93 | * | ||
| 94 | * These functions reduce a SIMD vector to a single scalar value by summing | ||
| 95 | * all elements. Used at the end of SIMD loops to extract the final result. | ||
| 96 | */ | ||
| 97 | |||
| 98 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 99 | |||
| 100 | #include <immintrin.h> | ||
| 101 | |||
| 102 | /* | ||
| 103 | * AVX-512 horizontal sum (16 floats -> 1 float) | ||
| 104 | * | ||
| 105 | * Uses _mm512_reduce_add_ps intrinsic (AVX-512F). | ||
| 106 | */ | ||
| 107 | VS_TARGET_AVX512 static inline float | ||
| 108 | 15967668 | vs_horizontal_sum_avx512(__m512 v) | |
| 109 | { | ||
| 110 | 15967668 | return _mm512_reduce_add_ps(v); | |
| 111 | } | ||
| 112 | |||
| 113 | /* | ||
| 114 | * AVX-512 horizontal sum for 64-bit integers (8 uint64_t -> 1 uint64_t) | ||
| 115 | * | ||
| 116 | * Used for Hamming distance (popcount results). | ||
| 117 | */ | ||
| 118 | VS_TARGET_AVX512 static inline uint64_t | ||
| 119 | 126 | vs_horizontal_sum_epi64_avx512(__m512i v) | |
| 120 | { | ||
| 121 | 126 | return _mm512_reduce_add_epi64(v); | |
| 122 | } | ||
| 123 | |||
| 124 | /* | ||
| 125 | * AVX2 horizontal sum (8 floats -> 1 float) | ||
| 126 | * | ||
| 127 | * AVX2 lacks a reduce intrinsic, so we manually combine the two 128-bit | ||
| 128 | * halves and use horizontal adds within each lane. | ||
| 129 | */ | ||
| 130 | VS_TARGET_AVX2 static inline float | ||
| 131 | 676 | vs_horizontal_sum_avx2(__m256 v) | |
| 132 | { | ||
| 133 | /* Extract high and low 128-bit halves */ | ||
| 134 | 676 | __m128 lo = _mm256_castps256_ps128(v); | |
| 135 | 676 | __m128 hi = _mm256_extractf128_ps(v, 1); | |
| 136 | |||
| 137 | /* Add halves */ | ||
| 138 | 676 | __m128 sum = _mm_add_ps(lo, hi); | |
| 139 | |||
| 140 | /* Horizontal add within 128 bits (twice to reduce to single element) */ | ||
| 141 | 676 | sum = _mm_hadd_ps(sum, sum); | |
| 142 | 676 | sum = _mm_hadd_ps(sum, sum); | |
| 143 | |||
| 144 | /* Extract scalar */ | ||
| 145 | 676 | return _mm_cvtss_f32(sum); | |
| 146 | } | ||
| 147 | |||
| 148 | /* | ||
| 149 | * AVX2 horizontal sum for 64-bit integers (4 uint64_t -> 1 uint64_t) | ||
| 150 | */ | ||
| 151 | VS_TARGET_AVX2 static inline uint64_t | ||
| 152 | 20 | vs_horizontal_sum_epi64_avx2(__m256i v) | |
| 153 | { | ||
| 154 | 20 | __m128i lo = _mm256_castsi256_si128(v); | |
| 155 | 20 | __m128i hi = _mm256_extracti128_si256(v, 1); | |
| 156 | 20 | __m128i sum = _mm_add_epi64(lo, hi); | |
| 157 | |||
| 158 | /* Extract two 64-bit values and add */ | ||
| 159 | 20 | uint64_t a = (uint64_t)_mm_extract_epi64(sum, 0); | |
| 160 | 20 | uint64_t b = (uint64_t)_mm_extract_epi64(sum, 1); | |
| 161 | 20 | return a + b; | |
| 162 | } | ||
| 163 | |||
| 164 | #endif /* x86_64 */ | ||
| 165 | |||
| 166 | #if defined(__aarch64__) || defined(_M_ARM64) | ||
| 167 | |||
| 168 | #include <arm_neon.h> | ||
| 169 | |||
| 170 | /* | ||
| 171 | * NEON horizontal sum (4 floats -> 1 float) | ||
| 172 | * | ||
| 173 | * Uses vaddvq_f32 (ARMv8.1+ reduction intrinsic). | ||
| 174 | */ | ||
| 175 | static inline float | ||
| 176 | vs_horizontal_sum_neon(float32x4_t v) | ||
| 177 | { | ||
| 178 | return vaddvq_f32(v); | ||
| 179 | } | ||
| 180 | |||
| 181 | /* | ||
| 182 | * NEON horizontal sum for 64-bit integers (2 uint64_t -> 1 uint64_t) | ||
| 183 | */ | ||
| 184 | static inline uint64_t | ||
| 185 | vs_horizontal_sum_u64_neon(uint64x2_t v) | ||
| 186 | { | ||
| 187 | return vaddvq_u64(v); | ||
| 188 | } | ||
| 189 | |||
| 190 | #endif /* aarch64 */ | ||
| 191 | |||
| 192 | /* | ||
| 193 | * Dispatch Macros | ||
| 194 | * | ||
| 195 | * These macros simplify runtime CPU detection and function pointer dispatch. | ||
| 196 | * Use DECLARE_DISPATCH to define the function pointer type and global state, | ||
| 197 | * then INIT_DISPATCH to initialize based on detected CPU capabilities. | ||
| 198 | * | ||
| 199 | * Example usage: | ||
| 200 | * | ||
| 201 | * // In header: | ||
| 202 | * Distance my_distance_fn(Vec32Ref a, Vec32Ref b); | ||
| 203 | * | ||
| 204 | * // In implementation: | ||
| 205 | * DECLARE_DISPATCH(my_distance, Distance, Vec32Ref, Vec32Ref) | ||
| 206 | * | ||
| 207 | * INIT_DISPATCH(my_distance, | ||
| 208 | * my_distance_scalar, | ||
| 209 | * my_distance_avx512, | ||
| 210 | * my_distance_avx2, | ||
| 211 | * my_distance_neon) | ||
| 212 | * | ||
| 213 | * Distance my_distance_fn(Vec32Ref a, Vec32Ref b) { | ||
| 214 | * if (vs_unlikely(!g_my_distance_initialized)) | ||
| 215 | * my_distance_init(); | ||
| 216 | * return g_my_distance_fn(a, b); | ||
| 217 | * } | ||
| 218 | */ | ||
| 219 | |||
| 220 | /* | ||
| 221 | * DECLARE_DISPATCH - Declare function pointer type and global state | ||
| 222 | * | ||
| 223 | * Creates: | ||
| 224 | * - typedef for the function pointer type | ||
| 225 | * - static global function pointer (initially NULL) | ||
| 226 | * - static atomic initialization flag | ||
| 227 | * | ||
| 228 | * The pointer is atomic because two threads racing the first call can | ||
| 229 | * both pass the flag check and both store to it. The stores pick the | ||
| 230 | * same implementation, but concurrent non-atomic stores to one object | ||
| 231 | * are a data race whatever they write. | ||
| 232 | */ | ||
| 233 | #define DECLARE_DISPATCH(name, ret_type, ...) \ | ||
| 234 | typedef ret_type (*name##_fn_t)(__VA_ARGS__); \ | ||
| 235 | static _Atomic(name##_fn_t) g_##name##_fn = NULL; \ | ||
| 236 | static _Atomic(bool) g_##name##_initialized = false; \ | ||
| 237 | static inline int name##_init(void); \ | ||
| 238 | static inline ret_type name##_dispatch(__VA_ARGS__); \ | ||
| 239 | \ | ||
| 240 | static inline ret_type name##_dispatch(__VA_ARGS__ args) \ | ||
| 241 | { \ | ||
| 242 | if (vs_unlikely(!g_##name##_initialized)) \ | ||
| 243 | name##_init(); \ | ||
| 244 | return g_##name##_fn(args); \ | ||
| 245 | } | ||
| 246 | |||
| 247 | /* | ||
| 248 | * INIT_DISPATCH - Initialize function pointer based on CPU capabilities | ||
| 249 | * | ||
| 250 | * Parameters: | ||
| 251 | * - name: base name (matches DECLARE_DISPATCH) | ||
| 252 | * - scalar: scalar fallback function | ||
| 253 | * - avx512: AVX-512 implementation (or NULL) | ||
| 254 | * - avx2: AVX2 implementation (or NULL) | ||
| 255 | * - neon: NEON implementation (or NULL) | ||
| 256 | * | ||
| 257 | * Selects the best available implementation based on runtime CPU detection. | ||
| 258 | * Falls back to scalar if no SIMD implementation is available. | ||
| 259 | */ | ||
| 260 | #define INIT_DISPATCH(name, scalar, avx512, avx2, neon) \ | ||
| 261 | static inline int name##_init(void) \ | ||
| 262 | { \ | ||
| 263 | if (g_##name##_initialized) \ | ||
| 264 | return 0; \ | ||
| 265 | \ | ||
| 266 | SimdCapability caps = vs_detect_simd(); \ | ||
| 267 | \ | ||
| 268 | /* Require every AVX-512 sub-extension any kernel \ | ||
| 269 | * here may use (F+DQ+BW), not just F -- see \ | ||
| 270 | * VS_SIMD_AVX512_* in platform.h. */ \ | ||
| 271 | SimdCapability avx512_req = VS_SIMD_AVX512_DQ | VS_SIMD_AVX512_BW; \ | ||
| 272 | if (((caps & avx512_req) == avx512_req) && (avx512) != NULL) \ | ||
| 273 | g_##name##_fn = avx512; \ | ||
| 274 | else if ((caps & SIMD_AVX2) && (avx2) != NULL) \ | ||
| 275 | g_##name##_fn = avx2; \ | ||
| 276 | else if ((caps & SIMD_NEON) && (neon) != NULL) \ | ||
| 277 | g_##name##_fn = neon; \ | ||
| 278 | else \ | ||
| 279 | g_##name##_fn = scalar; \ | ||
| 280 | \ | ||
| 281 | g_##name##_initialized = true; \ | ||
| 282 | return 0; \ | ||
| 283 | } | ||
| 284 | |||
| 285 | /* | ||
| 286 | * Batch Processing Helpers | ||
| 287 | * | ||
| 288 | * Common patterns for processing multiple vectors in a loop with prefetching. | ||
| 289 | */ | ||
| 290 | |||
| 291 | /* | ||
| 292 | * Prefetch distance (in vectors) for batch operations. | ||
| 293 | * | ||
| 294 | * Prefetch 2 vectors ahead to hide memory latency. Assumes 64-byte cache | ||
| 295 | * lines and typical vector sizes (128-1024 dimensions * 4 bytes = 0.5-4 KB). | ||
| 296 | */ | ||
| 297 | #define VS_PREFETCH_DISTANCE 2 | ||
| 298 | |||
| 299 | /* | ||
| 300 | * BATCH_LOOP_WITH_PREFETCH - Generic batch loop template | ||
| 301 | * | ||
| 302 | * Usage: | ||
| 303 | * BATCH_LOOP_WITH_PREFETCH(v, count, dim, vectors) { | ||
| 304 | * Vec32Ref vec = {.data = vectors + v * dim, .dim = dim}; | ||
| 305 | * distances[v] = compute_distance(query, vec); | ||
| 306 | * } | ||
| 307 | */ | ||
| 308 | #define BATCH_LOOP_WITH_PREFETCH(idx, count, dim, vectors) \ | ||
| 309 | for (uint32_t idx = 0; idx < (count); (idx)++) \ | ||
| 310 | { \ | ||
| 311 | if ((idx) + VS_PREFETCH_DISTANCE < (count)) \ | ||
| 312 | vs_prefetch_read( \ | ||
| 313 | (vectors) + ((idx) + VS_PREFETCH_DISTANCE) * (dim)); | ||
| 314 | |||
| 315 | #define BATCH_LOOP_END } | ||
| 316 | |||
| 317 | #endif /* VS_SIMD_UTILS_H */ | ||
| 318 |