Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion API.md
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,14 @@ SELECT vector_backend();
**Returns:** `TEXT`

**Description:**
Returns the active backend used by TurboQuant lookup-table scans. This is useful when validating that TurboQuant is using the expected SIMD path on a target runtime.
Returns the SIMD tier selected at load time, the same one `vector_backend()` reports.

TurboQuant lookup-table scans no longer vary by backend: the scan is one table lookup
per row, which is already about one load per cycle on any machine, and NEON has no
gather instruction at all. A single implementation is used everywhere, so the same query
returns the same distance whatever the CPU — the per-backend versions this replaced
differed by up to 1.5e-4 relative because they accumulated in `float` while the scalar
one accumulated in `double`.

**Example:**

Expand Down
50 changes: 2 additions & 48 deletions src/distance-avx2.c
Original file line number Diff line number Diff line change
Expand Up @@ -997,54 +997,7 @@ float bit1_distance_hamming_avx2 (const void *v1, const void *v2, int n) {
return (float)distance;
}

static inline uint16_t turbo_lut3_index_avx2 (const uint8_t *packed, int row, int packed_bytes) {
size_t bit_pos = (size_t)row * 12u;
size_t byte_pos = bit_pos / 8u;
int shift = (int)(bit_pos % 8u);
uint32_t word = 0;
if ((int)byte_pos < packed_bytes) word |= packed[byte_pos];
if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8;
return (uint16_t)((word >> shift) & 0x0fffu);
}

float turbo_lut_dot_avx2 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) {
__m256 acc = _mm256_setzero_ps();
const __m256i lane = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
int r = 0;
if (bits == 3) {
const __m256i stride = _mm256_set1_epi32(4096);
for (; r + 7 < lut_rows; r += 8) {
int idx[8];
for (int i = 0; i < 8; ++i) idx[i] = turbo_lut3_index_avx2(packed, r + i, packed_bytes);
__m256i codes = _mm256_loadu_si256((const __m256i *)idx);
__m256i rows = _mm256_add_epi32(_mm256_set1_epi32(r), lane);
__m256i indices = _mm256_add_epi32(_mm256_mullo_epi32(rows, stride), codes);
__m256 vals = _mm256_i32gather_ps(query_lut, indices, 4);
acc = _mm256_add_ps(acc, vals);
}
} else {
const __m256i stride = _mm256_set1_epi32(256);
for (; r + 7 < lut_rows; r += 8) {
__m128i codes8 = _mm_loadl_epi64((const __m128i *)(packed + r));
__m256i codes = _mm256_cvtepu8_epi32(codes8);
__m256i rows = _mm256_add_epi32(_mm256_set1_epi32(r), lane);
__m256i indices = _mm256_add_epi32(_mm256_mullo_epi32(rows, stride), codes);
__m256 vals = _mm256_i32gather_ps(query_lut, indices, 4);
acc = _mm256_add_ps(acc, vals);
}
}

float partial[8];
_mm256_storeu_ps(partial, acc);
float dot = partial[0] + partial[1] + partial[2] + partial[3] +
partial[4] + partial[5] + partial[6] + partial[7];
if (bits == 3) {
for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_avx2(packed, r, packed_bytes)];
} else {
for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]];
}
return dot * scale;
}

#endif

Expand Down Expand Up @@ -1085,7 +1038,8 @@ bool init_distance_functions_avx2 (void) {
dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_avx2;

distance_backend_name = "AVX2";
turbo_lut_dot_function = turbo_lut_dot_avx2;
// the TurboQuant lookup scan is gather-bound and shared by every backend
turbo_lut_dot_function = turbo_lut_dot_cpu;
turbo_lut_backend_name = "AVX2";
return true;
#else
Expand Down
1 change: 0 additions & 1 deletion src/distance-avx2.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,5 @@

// returns true when the AVX2 kernels were compiled into this build
bool init_distance_functions_avx2 (void);
float turbo_lut_dot_avx2 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes);

#endif
47 changes: 2 additions & 45 deletions src/distance-avx512.c
Original file line number Diff line number Diff line change
Expand Up @@ -976,51 +976,7 @@ static float bit1_distance_hamming_avx512(const void *v1, const void *v2, int n)
return (float)distance;
}

static inline uint16_t turbo_lut3_index_avx512 (const uint8_t *packed, int row, int packed_bytes) {
size_t bit_pos = (size_t)row * 12u;
size_t byte_pos = bit_pos / 8u;
int shift = (int)(bit_pos % 8u);
uint32_t word = 0;
if ((int)byte_pos < packed_bytes) word |= packed[byte_pos];
if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8;
return (uint16_t)((word >> shift) & 0x0fffu);
}

float turbo_lut_dot_avx512 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) {
__m512 acc = _mm512_setzero_ps();
const __m512i lane = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
int r = 0;
if (bits == 3) {
const __m512i stride = _mm512_set1_epi32(4096);
for (; r + 15 < lut_rows; r += 16) {
int idx[16];
for (int i = 0; i < 16; ++i) idx[i] = turbo_lut3_index_avx512(packed, r + i, packed_bytes);
__m512i codes = _mm512_loadu_si512((const void *)idx);
__m512i rows = _mm512_add_epi32(_mm512_set1_epi32(r), lane);
__m512i indices = _mm512_add_epi32(_mm512_mullo_epi32(rows, stride), codes);
__m512 vals = _mm512_i32gather_ps(indices, query_lut, 4);
acc = _mm512_add_ps(acc, vals);
}
} else {
const __m512i stride = _mm512_set1_epi32(256);
for (; r + 15 < lut_rows; r += 16) {
__m128i codes8 = _mm_loadu_si128((const __m128i *)(packed + r));
__m512i codes = _mm512_cvtepu8_epi32(codes8);
__m512i rows = _mm512_add_epi32(_mm512_set1_epi32(r), lane);
__m512i indices = _mm512_add_epi32(_mm512_mullo_epi32(rows, stride), codes);
__m512 vals = _mm512_i32gather_ps(indices, query_lut, 4);
acc = _mm512_add_ps(acc, vals);
}
}

float dot = _mm512_reduce_add_ps(acc);
if (bits == 3) {
for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_avx512(packed, r, packed_bytes)];
} else {
for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]];
}
return dot * scale;
}

#endif

Expand Down Expand Up @@ -1061,7 +1017,8 @@ bool init_distance_functions_avx512(void) {
dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_avx512;

distance_backend_name = "AVX512";
turbo_lut_dot_function = turbo_lut_dot_avx512;
// the TurboQuant lookup scan is gather-bound and shared by every backend
turbo_lut_dot_function = turbo_lut_dot_cpu;
turbo_lut_backend_name = "AVX512";
return true;
#else
Expand Down
1 change: 0 additions & 1 deletion src/distance-avx512.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,5 @@

// returns true when the AVX512 kernels were compiled into this build
bool init_distance_functions_avx512 (void);
float turbo_lut_dot_avx512 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes);

#endif
35 changes: 29 additions & 6 deletions src/distance-cpu.c
Original file line number Diff line number Diff line change
Expand Up @@ -879,19 +879,42 @@ static inline uint16_t turbo_lut3_index_cpu (const uint8_t *packed, int row, int
return (uint16_t)((word >> shift) & 0x0fffu);
}

// The TurboQuant scan is a chain of table lookups: one gather per row, and on any
// machine that is already about one load per cycle. There is nothing for SIMD to do -
// NEON has no gather at all, and the four per-backend copies this replaces were scalar
// gathers into a stack array plus a single vector add. What they were really buying was
// four parallel float lanes instead of one serial double accumulator, and four
// independent double accumulators buy the same parallelism without the accuracy loss:
// measured within 2% of the NEON version at every bit width, and identical on every
// backend rather than differing by up to 1.5e-4 relative depending on which one ran.
float turbo_lut_dot_cpu (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) {
double dot = 0.0;
double acc0 = 0.0, acc1 = 0.0, acc2 = 0.0, acc3 = 0.0;
int r = 0;

if (bits == 3) {
for (int r = 0; r < lut_rows; ++r) {
dot += (double)query_lut[(size_t)r * 4096u + turbo_lut3_index_cpu(packed, r, packed_bytes)];
for (; r + 3 < lut_rows; r += 4) {
acc0 += (double)query_lut[(size_t)(r + 0) * 4096u + turbo_lut3_index_cpu(packed, r + 0, packed_bytes)];
acc1 += (double)query_lut[(size_t)(r + 1) * 4096u + turbo_lut3_index_cpu(packed, r + 1, packed_bytes)];
acc2 += (double)query_lut[(size_t)(r + 2) * 4096u + turbo_lut3_index_cpu(packed, r + 2, packed_bytes)];
acc3 += (double)query_lut[(size_t)(r + 3) * 4096u + turbo_lut3_index_cpu(packed, r + 3, packed_bytes)];
}
for (; r < lut_rows; ++r) {
acc0 += (double)query_lut[(size_t)r * 4096u + turbo_lut3_index_cpu(packed, r, packed_bytes)];
}
} else {
(void)packed_bytes;
for (int r = 0; r < lut_rows; ++r) {
dot += (double)query_lut[(size_t)r * 256u + packed[r]];
for (; r + 3 < lut_rows; r += 4) {
acc0 += (double)query_lut[(size_t)(r + 0) * 256u + packed[r + 0]];
acc1 += (double)query_lut[(size_t)(r + 1) * 256u + packed[r + 1]];
acc2 += (double)query_lut[(size_t)(r + 2) * 256u + packed[r + 2]];
acc3 += (double)query_lut[(size_t)(r + 3) * 256u + packed[r + 3]];
}
for (; r < lut_rows; ++r) {
acc0 += (double)query_lut[(size_t)r * 256u + packed[r]];
}
}
return (float)(dot * (double)scale);

return (float)(((acc0 + acc1) + (acc2 + acc3)) * (double)scale);
}

void init_cpu_functions (void) {
Expand Down
1 change: 1 addition & 0 deletions src/distance-cpu.h
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ typedef float (*turbo_lut_dot_function_t)(const uint8_t *packed, float scale, co
void init_distance_functions (bool force_cpu);

extern turbo_lut_dot_function_t turbo_lut_dot_function;
float turbo_lut_dot_cpu (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes);
extern const char *turbo_lut_backend_name;

// MARK: - FLOAT16/BFLOAT16 -
Expand Down
52 changes: 2 additions & 50 deletions src/distance-neon.c
Original file line number Diff line number Diff line change
Expand Up @@ -1200,56 +1200,7 @@ float bit1_distance_hamming_neon (const void *v1, const void *v2, int n) {
return (float)distance;
}

static inline uint16_t turbo_lut3_index_neon (const uint8_t *packed, int row, int packed_bytes) {
size_t bit_pos = (size_t)row * 12u;
size_t byte_pos = bit_pos / 8u;
int shift = (int)(bit_pos % 8u);
uint32_t word = 0;
if ((int)byte_pos < packed_bytes) word |= packed[byte_pos];
if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8;
return (uint16_t)((word >> shift) & 0x0fffu);
}

float turbo_lut_dot_neon (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) {
float32x4_t acc = vdupq_n_f32(0.0f);
int r = 0;
if (bits == 3) {
for (; r + 3 < lut_rows; r += 4) {
float tmp[4] = {
query_lut[(size_t)(r + 0) * 4096u + turbo_lut3_index_neon(packed, r + 0, packed_bytes)],
query_lut[(size_t)(r + 1) * 4096u + turbo_lut3_index_neon(packed, r + 1, packed_bytes)],
query_lut[(size_t)(r + 2) * 4096u + turbo_lut3_index_neon(packed, r + 2, packed_bytes)],
query_lut[(size_t)(r + 3) * 4096u + turbo_lut3_index_neon(packed, r + 3, packed_bytes)]
};
acc = vaddq_f32(acc, vld1q_f32(tmp));
}
} else {
for (; r + 3 < lut_rows; r += 4) {
float tmp[4] = {
query_lut[(size_t)(r + 0) * 256u + packed[r + 0]],
query_lut[(size_t)(r + 1) * 256u + packed[r + 1]],
query_lut[(size_t)(r + 2) * 256u + packed[r + 2]],
query_lut[(size_t)(r + 3) * 256u + packed[r + 3]]
};
acc = vaddq_f32(acc, vld1q_f32(tmp));
}
}

float dot;
#if defined(__aarch64__)
dot = vaddvq_f32(acc);
#else
float partial[4];
vst1q_f32(partial, acc);
dot = partial[0] + partial[1] + partial[2] + partial[3];
#endif
if (bits == 3) {
for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_neon(packed, r, packed_bytes)];
} else {
for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]];
}
return dot * scale;
}

#endif

Expand Down Expand Up @@ -1290,7 +1241,8 @@ bool init_distance_functions_neon (void) {
dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_neon;

distance_backend_name = "NEON";
turbo_lut_dot_function = turbo_lut_dot_neon;
// the TurboQuant lookup scan is gather-bound and shared by every backend
turbo_lut_dot_function = turbo_lut_dot_cpu;
turbo_lut_backend_name = "NEON";
return true;
#else
Expand Down
1 change: 0 additions & 1 deletion src/distance-neon.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,5 @@

// returns true when the NEON kernels were compiled into this build
bool init_distance_functions_neon (void);
float turbo_lut_dot_neon (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes);

#endif
31 changes: 2 additions & 29 deletions src/distance-rvv.c
Original file line number Diff line number Diff line change
Expand Up @@ -988,35 +988,7 @@ float bit1_distance_hamming_rvv (const void *v1, const void *v2, int n) {
return (float) uint64_sum_vector_u64m8(vdistance, vl);
}

static inline uint16_t turbo_lut3_index_rvv (const uint8_t *packed, int row, int packed_bytes) {
size_t bit_pos = (size_t)row * 12u;
size_t byte_pos = bit_pos / 8u;
int shift = (int)(bit_pos % 8u);
uint32_t word = 0;
if ((int)byte_pos < packed_bytes) word |= packed[byte_pos];
if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8;
return (uint16_t)((word >> shift) & 0x0fffu);
}

float turbo_lut_dot_rvv (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) {
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
int r = 0;
while (r < lut_rows) {
size_t n = (size_t)(lut_rows - r);
size_t vl = __riscv_vsetvl_e32m8(n);
float tmp[vl];
for (size_t i = 0; i < vl; ++i) {
int row = r + (int)i;
if (bits == 3) tmp[i] = query_lut[(size_t)row * 4096u + turbo_lut3_index_rvv(packed, row, packed_bytes)];
else tmp[i] = query_lut[(size_t)row * 256u + packed[row]];
}
vfloat32m8_t vals = __riscv_vle32_v_f32m8(tmp, vl);
acc = __riscv_vfadd_vv_f32m8(acc, vals, vl);
r += (int)vl;
}
return float32_sum_vector_f32m8(acc, vlmax) * scale;
}
#endif

// MARK: -
Expand Down Expand Up @@ -1056,7 +1028,8 @@ bool init_distance_functions_rvv (void) {
dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_rvv;

distance_backend_name = "RVV";
turbo_lut_dot_function = turbo_lut_dot_rvv;
// the TurboQuant lookup scan is gather-bound and shared by every backend
turbo_lut_dot_function = turbo_lut_dot_cpu;
turbo_lut_backend_name = "RVV";
return true;
#else
Expand Down
1 change: 0 additions & 1 deletion src/distance-rvv.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,5 @@

// returns true when the RVV kernels were compiled into this build
bool init_distance_functions_rvv (void);
float turbo_lut_dot_rvv (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes);

#endif
Loading
Loading