mirror of
https://github.com/redis/redis.git
synced 2026-02-03 20:39:54 -05:00
37f685908e
Some checks are pending
CI / test-ubuntu-latest (push) Waiting to run
CI / test-sanitizer-address (push) Waiting to run
CI / build-debian-old (push) Waiting to run
CI / build-macos-latest (push) Waiting to run
CI / build-32bit (push) Waiting to run
CI / build-libc-malloc (push) Waiting to run
CI / build-centos-jemalloc (push) Waiting to run
CI / build-old-chain-jemalloc (push) Waiting to run
Codecov / code-coverage (push) Waiting to run
External Server Tests / test-external-standalone (push) Waiting to run
External Server Tests / test-external-cluster (push) Waiting to run
External Server Tests / test-external-nodebug (push) Waiting to run
Spellcheck / Spellcheck (push) Waiting to run
This pull request vectorizes the 8-bit quantization vector-search path in a similar was as the non-quantization path. The assembly intrinsics are a bit more complicated than in the non-quantization path, since we are operating on 8-bit integers and we need to worry about preventing overflow. Thus, after loading the 8-bit integers, they are extended into 16-bits before multiplying and accumulating into 32-bit integers. --------- Co-authored-by: debing.sun <debing.sun@redis.com>
85 lines
4.6 KiB
Python
85 lines
4.6 KiB
Python
from test import TestCase
|
|
|
|
class Q8Vectorization(TestCase):
|
|
def getname(self):
|
|
return "Q8 quantization: verify vectorized vs scalar paths produce consistent results"
|
|
|
|
def test(self):
|
|
# Test with different dimensions to exercise different code paths and boundaries:
|
|
# - dim=16: Scalar path (< 32)
|
|
# - dim=31: Largest scalar-only dimension (boundary)
|
|
# - dim=32: Smallest AVX2 dimension, no remainder (boundary)
|
|
# - dim=33: AVX2 with 1-element remainder
|
|
# - dim=63: AVX2 with 31-element remainder (largest AVX2-only)
|
|
# - dim=64: Smallest AVX512 dimension, no remainder (boundary)
|
|
# - dim=65: AVX512 with 1-element remainder
|
|
# - dim=128: AVX512 path with no remainder
|
|
# - dim=256, dim=512: Large dimensions to test overflow prevention
|
|
|
|
test_dims = [16, 31, 32, 33, 63, 64, 65, 128, 256, 512]
|
|
|
|
for dim in test_dims:
|
|
key = f'{self.test_key}:dim{dim}'
|
|
|
|
# Test vectors with extreme values to verify overflow prevention:
|
|
# vec1: all +1.0 -> quantizes to +127 (max positive int8)
|
|
# vec2: all +0.99 -> quantizes to ~+126 (similar to vec1)
|
|
# vec3: all -1.0 -> quantizes to -127/-128 (max negative int8)
|
|
# vec4: alternating +1.0/-1.0 -> alternating +127/-127 (tests mixed signs)
|
|
vec1 = [1.0] * dim # All max positive
|
|
vec2 = [0.99] * dim # Similar to vec1
|
|
vec3 = [-1.0] * dim # All max negative (opposite direction)
|
|
vec4 = [1.0 if i % 2 == 0 else -1.0 for i in range(dim)] # Alternating extreme values
|
|
|
|
# Add vectors with Q8 quantization
|
|
self.redis.execute_command('VADD', key, 'VALUES', dim,
|
|
*[str(x) for x in vec1], f'{key}:item:1', 'Q8')
|
|
self.redis.execute_command('VADD', key, 'VALUES', dim,
|
|
*[str(x) for x in vec2], f'{key}:item:2', 'Q8')
|
|
self.redis.execute_command('VADD', key, 'VALUES', dim,
|
|
*[str(x) for x in vec3], f'{key}:item:3', 'Q8')
|
|
self.redis.execute_command('VADD', key, 'VALUES', dim,
|
|
*[str(x) for x in vec4], f'{key}:item:4', 'Q8')
|
|
|
|
# Query similarity using vec1 (all max positive values)
|
|
# This exercises worst-case positive accumulation: dim * 127 * 127
|
|
result = self.redis.execute_command('VSIM', key, 'VALUES', dim,
|
|
*[str(x) for x in vec1], 'WITHSCORES')
|
|
|
|
# Convert results to dictionary
|
|
results_dict = {}
|
|
for i in range(0, len(result), 2):
|
|
k = result[i].decode()
|
|
score = float(result[i+1])
|
|
results_dict[k] = score
|
|
|
|
# Verify results - these would be wrong if overflow occurred
|
|
# Self-similarity should be ~1.0 (identical vectors)
|
|
assert results_dict[f'{key}:item:1'] > 0.99, \
|
|
f"Dim {dim}: Self-similarity too low: {results_dict[f'{key}:item:1']}"
|
|
|
|
# Similar vector should have high similarity
|
|
assert results_dict[f'{key}:item:2'] > 0.99, \
|
|
f"Dim {dim}: Similar vector similarity too low: {results_dict[f'{key}:item:2']}"
|
|
|
|
# Opposite vector should have very low similarity (~0.0)
|
|
# With overflow bug, this could give incorrect positive values
|
|
assert results_dict[f'{key}:item:3'] < 0.1, \
|
|
f"Dim {dim}: Opposite vector similarity too high: {results_dict[f'{key}:item:3']}"
|
|
|
|
# Alternating vector: dot product sums to ~0, so similarity ~0.5
|
|
# (127*127) + (127*-127) + ... = 0, normalized gives ~0.5
|
|
assert 0.4 < results_dict[f'{key}:item:4'] < 0.6, \
|
|
f"Dim {dim}: Alternating vector similarity unexpected: {results_dict[f'{key}:item:4']}"
|
|
|
|
# Also query with the alternating pattern to verify its self-similarity
|
|
result_alt = self.redis.execute_command('VSIM', key, 'VALUES', dim,
|
|
*[str(x) for x in vec4], 'WITHSCORES')
|
|
results_alt = {}
|
|
for i in range(0, len(result_alt), 2):
|
|
k = result_alt[i].decode()
|
|
score = float(result_alt[i+1])
|
|
results_alt[k] = score
|
|
|
|
assert results_alt[f'{key}:item:4'] > 0.99, \
|
|
f"Dim {dim}: Alternating self-similarity too low: {results_alt[f'{key}:item:4']}"
|