| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * fastscan_avx2.c - AVX2 VPSHUFB accumulate kernel | ||
| 6 | * | ||
| 7 | * AVX2 version of the even/odd byte accumulation. Processes 64B | ||
| 8 | * per iteration in two 32B loads (since AVX2 registers are 256-bit). | ||
| 9 | * Each iteration handles 4 subquantizers. | ||
| 10 | */ | ||
| 11 | |||
| 12 | #include "vs_config.h" | ||
| 13 | |||
| 14 | #ifdef VS_SIMD_FULL | ||
| 15 | |||
| 16 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 17 | |||
| 18 | #include <immintrin.h> | ||
| 19 | |||
| 20 | #include "algo/simd_utils.h" | ||
| 21 | #include "quant/fastscan.h" | ||
| 22 | |||
| 23 | /* Accumulate one 32-byte block into even/odd accumulators */ | ||
| 24 | #define ACCUM_32B(codes_ptr, lut_ptr, a0, a1, a2, a3, mask) \ | ||
| 25 | do \ | ||
| 26 | { \ | ||
| 27 | __m256i c_ = _mm256_loadu_si256((const __m256i *)(codes_ptr)); \ | ||
| 28 | __m256i tab_ = _mm256_loadu_si256((const __m256i *)(lut_ptr)); \ | ||
| 29 | __m256i lo_ = _mm256_and_si256(c_, mask); \ | ||
| 30 | __m256i hi_ = _mm256_and_si256(_mm256_srli_epi16(c_, 4), mask); \ | ||
| 31 | __m256i rlo_ = _mm256_shuffle_epi8(tab_, lo_); \ | ||
| 32 | __m256i rhi_ = _mm256_shuffle_epi8(tab_, hi_); \ | ||
| 33 | a0 = _mm256_add_epi16(a0, rlo_); \ | ||
| 34 | a1 = _mm256_add_epi16(a1, _mm256_srli_epi16(rlo_, 8)); \ | ||
| 35 | a2 = _mm256_add_epi16(a2, rhi_); \ | ||
| 36 | a3 = _mm256_add_epi16(a3, _mm256_srli_epi16(rhi_, 8)); \ | ||
| 37 | } while (0) | ||
| 38 | |||
| 39 | /* Reduce even/odd accumulators to 16 uint16 results */ | ||
| 40 | static inline VS_TARGET_AVX2 __m256i | ||
| 41 | 7539844 | reduce_accu_pair(__m256i even, __m256i odd) | |
| 42 | { | ||
| 43 | 7539844 | even = _mm256_sub_epi16(even, _mm256_slli_epi16(odd, 8)); | |
| 44 | 7539844 | return _mm256_add_epi16( | |
| 45 | 289344 | _mm256_permute2f128_si256(even, odd, 0x21), | |
| 46 | 289344 | _mm256_blend_epi32(even, odd, 0xF0)); | |
| 47 | } | ||
| 48 | |||
| 49 | VS_TARGET_AVX2 void | ||
| 50 | 48 | vs_fastscan_accumulate_avx2( | |
| 51 | const uint8_t *codes, | ||
| 52 | const uint8_t *lut, | ||
| 53 | uint16_t *accum, | ||
| 54 | Dimension dim) | ||
| 55 | { | ||
| 56 | 48 | uint32_t code_length = VS_FASTSCAN_GROUP_BYTES(dim); | |
| 57 | |||
| 58 | 48 | __m256i low_mask = _mm256_set1_epi8(0x0F); | |
| 59 | 48 | __m256i accu0 = _mm256_setzero_si256(); | |
| 60 | 48 | __m256i accu1 = _mm256_setzero_si256(); | |
| 61 | 48 | __m256i accu2 = _mm256_setzero_si256(); | |
| 62 | 48 | __m256i accu3 = _mm256_setzero_si256(); | |
| 63 | |||
| 64 | /* One 32B column (8 dims) per iteration. code_length is a multiple | ||
| 65 | * of 32 but not necessarily 64, so stepping 64B (two columns) at a | ||
| 66 | * time would read a phantom column past the code region whenever | ||
| 67 | * ceil(dim/8) is odd. The accumulators are order-independent, so a | ||
| 68 | * 32B stride is equivalent and never over-reads. */ | ||
| 69 |
2/2✓ Branch 0 taken 3052 times.
✓ Branch 1 taken 48 times.
|
3100 | for (uint32_t i = 0; i < code_length; i += 32) |
| 70 | 42728 | ACCUM_32B(codes + i, lut + i, accu0, accu1, accu2, accu3, low_mask); | |
| 71 | |||
| 72 | 48 | _mm256_storeu_si256((__m256i *)accum, reduce_accu_pair(accu0, accu1)); | |
| 73 | 48 | _mm256_storeu_si256( | |
| 74 | 48 | (__m256i *)(accum + 16), reduce_accu_pair(accu2, accu3)); | |
| 75 | 48 | } | |
| 76 | |||
| 77 | #undef ACCUM_32B | ||
| 78 | |||
| 79 | /* ---------------------------------------------------------------- | ||
| 80 | * AVX2 high-accuracy accumulate (uint16 LUT → int32 output) | ||
| 81 | * | ||
| 82 | * Two VPSHUFB passes per 32B code block (lo-byte table, hi-byte | ||
| 83 | * table). 8 accumulators: accu[lo/hi][0-3]. | ||
| 84 | * ---------------------------------------------------------------- */ | ||
| 85 | |||
| 86 | #define HACC_ACCUM_256(tab, lo, hi, a) \ | ||
| 87 | do \ | ||
| 88 | { \ | ||
| 89 | __m256i rlo_ = _mm256_shuffle_epi8(tab, lo); \ | ||
| 90 | __m256i rhi_ = _mm256_shuffle_epi8(tab, hi); \ | ||
| 91 | a[0] = _mm256_add_epi16(a[0], rlo_); \ | ||
| 92 | a[1] = _mm256_add_epi16(a[1], _mm256_srli_epi16(rlo_, 8)); \ | ||
| 93 | a[2] = _mm256_add_epi16(a[2], rhi_); \ | ||
| 94 | a[3] = _mm256_add_epi16(a[3], _mm256_srli_epi16(rhi_, 8)); \ | ||
| 95 | } while (0) | ||
| 96 | |||
| 97 | /* Reduce even/odd accumulators to 16 uint16, then widen to int32. | ||
| 98 | * Returns two __m256i: lo 8 int32, hi 8 int32 via output pointers. */ | ||
| 99 | static inline VS_TARGET_AVX2 void | ||
| 100 | 7539748 | hacc_reduce_avx2(__m256i even, __m256i odd, __m256i *out_lo8, __m256i *out_hi8) | |
| 101 | { | ||
| 102 | 7539748 | __m256i r16 = reduce_accu_pair(even, odd); | |
| 103 | 7539748 | *out_lo8 = _mm256_cvtepu16_epi32(_mm256_castsi256_si128(r16)); | |
| 104 | 7539748 | *out_hi8 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256(r16, 1)); | |
| 105 | 289248 | } | |
| 106 | |||
| 107 | VS_TARGET_AVX2 void | ||
| 108 | 3697562 | vs_fastscan_accumulate_hacc_avx2( | |
| 109 | const uint8_t *codes, | ||
| 110 | const uint8_t *lut, | ||
| 111 | int32_t *accum, | ||
| 112 | Dimension dim) | ||
| 113 | { | ||
| 114 | 3697562 | __m256i low_mask = _mm256_set1_epi8(0x0F); | |
| 115 | |||
| 116 | 3625250 | __m256i accu[2][4]; | |
| 117 |
2/2✓ Branch 0 taken 7395124 times.
✓ Branch 1 taken 3697562 times.
|
11092686 | for (int q = 0; q < 2; q++) |
| 118 |
2/2✓ Branch 0 taken 29580496 times.
✓ Branch 1 taken 7395124 times.
|
36975620 | for (int r = 0; r < 4; r++) |
| 119 | 29580496 | accu[q][r] = _mm256_setzero_si256(); | |
| 120 | |||
| 121 | /* One 32B column (2 subquantizers) per iteration. The high-accuracy | ||
| 122 | * LUT is laid out in 128B blocks of 4 subquantizers ([4x16 lo][4x16 | ||
| 123 | * hi]); column c lives in block c/2 at within-block offset (c%2)*32, | ||
| 124 | * with the hi table 64B after the lo. Iterating whole columns reads | ||
| 125 | * exactly code_length bytes -- stepping 64B (4 sq) at a time would | ||
| 126 | * read a phantom column past the code region when ceil(dim/8) is | ||
| 127 | * odd. Trailing phantom subquantizers (nsq not a multiple of 2) have | ||
| 128 | * a zeroed LUT, so they contribute nothing. */ | ||
| 129 | 3697562 | uint32_t ncols = VS_FASTSCAN_GROUP_BYTES(dim) / VS_FASTSCAN_GROUP; | |
| 130 |
2/2✓ Branch 0 taken 30216433 times.
✓ Branch 1 taken 3697562 times.
|
33913995 | for (uint32_t c = 0; c < ncols; c++) |
| 131 | { | ||
| 132 | 30417785 | __m256i cc = _mm256_loadu_si256((const __m256i *)(codes + c * 32)); | |
| 133 | 30216433 | __m256i lo = _mm256_and_si256(cc, low_mask); | |
| 134 | 30216433 | __m256i hi = _mm256_and_si256(_mm256_srli_epi16(cc, 4), low_mask); | |
| 135 | |||
| 136 | 30216433 | const uint8_t *lut_blk = lut + (c / 2) * 128 + (c % 2) * 32; | |
| 137 | 30216433 | __m256i tab_lo = _mm256_loadu_si256((const __m256i *)lut_blk); | |
| 138 | 31424545 | HACC_ACCUM_256(tab_lo, lo, hi, accu[0]); | |
| 139 | |||
| 140 | 30417785 | __m256i tab_hi = _mm256_loadu_si256((const __m256i *)(lut_blk + 64)); | |
| 141 | 31625897 | HACC_ACCUM_256(tab_hi, lo, hi, accu[1]); | |
| 142 | } | ||
| 143 | |||
| 144 | /* Reduce: vectors 0-15 (accu[][0,1]), 16-31 (accu[][2,3]) */ | ||
| 145 |
2/2✓ Branch 0 taken 7395124 times.
✓ Branch 1 taken 3697562 times.
|
11092686 | for (int half = 0; half < 2; half++) |
| 146 | { | ||
| 147 | 7250500 | __m256i lo_a, lo_b, hi_a, hi_b; | |
| 148 | 7395124 | hacc_reduce_avx2( | |
| 149 | 7395124 | accu[0][half * 2], accu[0][half * 2 + 1], &lo_a, &lo_b); | |
| 150 | 7395124 | hacc_reduce_avx2( | |
| 151 | 144624 | accu[1][half * 2], accu[1][half * 2 + 1], &hi_a, &hi_b); | |
| 152 | 7395124 | int32_t *out = accum + half * 16; | |
| 153 | 7684372 | _mm256_storeu_si256( | |
| 154 | (__m256i *)out, | ||
| 155 | _mm256_add_epi32(lo_a, _mm256_slli_epi32(hi_a, 8))); | ||
| 156 | 7539748 | _mm256_storeu_si256( | |
| 157 | 7395124 | (__m256i *)(out + 8), | |
| 158 | _mm256_add_epi32(lo_b, _mm256_slli_epi32(hi_b, 8))); | ||
| 159 | } | ||
| 160 | 3697562 | } | |
| 161 | |||
| 162 | #undef HACC_ACCUM_256 | ||
| 163 | |||
| 164 | #endif /* x86_64 */ | ||
| 165 | |||
| 166 | #endif /* VS_SIMD_FULL */ | ||
| 167 |