upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-08-06 18:10:58 -05:00
Code Issues Packages Projects Releases Wiki Activity

horizontal sum and test

Browse Source
This commit is contained in:
Ivan Pleshkov
2022-01-04 16:30:36 +03:00
parent 70e376bf5e
commit cd2f334271
1 changed files with 72 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
lib/segment/src/spaces/simple.rs
View File

@@ -14,6 +14,19 @@ pub struct CosineMetric {}
pub struct EuclidMetric {}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn hsum256_ps_avx(x: __m256) -> f32 {
/* ( x3+x7, x2+x6, x1+x5, x0+x4 ) */
let x128 : __m128 = _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x));
/* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */
let x64 : __m128 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128));
/* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */
let x32 : __m128 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55));
/* Conversion to float is a no-op on x86-64 */
return _mm_cvtss_f32(x32);
}
fn euclid_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
let s: ScoreType = v1
.iter()
@@ -25,19 +38,19 @@ fn euclid_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> Scor
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn euclid_similarity_avx2(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
unsafe {
let mut sum256: __m256 = _mm256_setzero_ps();
for i in (0..n).step_by(8) {
let sub256: __m256 = _mm256_sub_ps(_mm256_loadu_ps(&v1[i]), _mm256_loadu_ps(&v2[i]));
sum256 = _mm256_fmadd_ps(sub256, sub256, sum256);
}
let res: f32 = hsum256_ps_avx(sum256);
for i in (n - (n % 8)..n) {
res += (v1[i] - v2[i]).powi(2);
}
-res.sqrt()
#[target_feature(enable = "avx2")]
unsafe fn euclid_similarity_avx2(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
let n = v1.len();
let mut sum256: __m256 = _mm256_setzero_ps();
for i in (0..n).step_by(8) {
let sub256: __m256 = _mm256_sub_ps(_mm256_loadu_ps(&v1[i]), _mm256_loadu_ps(&v2[i]));
sum256 = _mm256_fmadd_ps(sub256, sub256, sum256);
}
let mut res = hsum256_ps_avx(sum256);
for i in n - (n % 8)..n {
res += (v1[i] - v2[i]).powi(2);
}
-res.sqrt()
}
fn dot_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
@@ -45,18 +58,18 @@ fn dot_similarity(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreTy
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn dot_similarity_avx2(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
unsafe {
let mut sum256: __m256 = _mm256_setzero_ps();
for i in (0..n).step_by(8) {
sum256 = _mm256_fmadd_ps(_mm256_loadu_ps(&v1[i]), _mm256_loadu_ps(&v2[i]), sum256);
}
let res: f32 = hsum256_ps_avx(sum256);
for i in (n - (n % 8)..n) {
res += v1[i] * v2[i];
}
res
#[target_feature(enable = "avx2")]
unsafe fn dot_similarity_avx2(v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
let n = v1.len();
let mut sum256: __m256 = _mm256_setzero_ps();
for i in (0..n).step_by(8) {
sum256 = _mm256_fmadd_ps(_mm256_loadu_ps(&v1[i]), _mm256_loadu_ps(&v2[i]), sum256);
}
let mut res = hsum256_ps_avx(sum256);
for i in n - (n % 8)..n {
res += v1[i] * v2[i];
}
res
}
fn cosine_preprocess(vector: &[VectorElementType]) -> Vec<VectorElementType> {
@@ -66,20 +79,20 @@ fn cosine_preprocess(vector: &[VectorElementType]) -> Vec<VectorElementType> {
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn cosine_preprocess_avx2(vector: &[VectorElementType]) -> Vec<VectorElementType> {
let length = unsafe {
let mut sum256: __m256 = _mm256_setzero_ps();
for i in (0..n).step_by(8) {
sum256 = _mm256_fmadd_ps(
_mm256_loadu_ps(&vector[i]),
_mm256_loadu_ps(&vector[i]), sum256);
}
let res: f32 = hsum256_ps_avx(sum256);
for i in (n - (n % 8)..n) {
res += vector[i].powi(2);
}
res
}.sqrt();
#[target_feature(enable = "avx2")]
unsafe fn cosine_preprocess_avx2(vector: &[VectorElementType]) -> Vec<VectorElementType> {
let n = vector.len();
let mut sum256: __m256 = _mm256_setzero_ps();
for i in (0..n).step_by(8) {
sum256 = _mm256_fmadd_ps(
_mm256_loadu_ps(&vector[i]),
_mm256_loadu_ps(&vector[i]), sum256);
}
let mut length = hsum256_ps_avx(sum256);
for i in n - (n % 8)..n {
length += vector[i].powi(2);
}
length = length.sqrt();
vector.iter().map(|x| x / length).collect()
}
@@ -91,8 +104,8 @@ impl Metric for EuclidMetric {
fn similarity(&self, v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if is_x86_feature_detected!("avx2") == 0 {
return euclid_similarity_avx2(v1, v2);
if is_x86_feature_detected!("avx2") {
return unsafe { euclid_similarity_avx2(v1, v2) };
}
}
euclid_similarity(v1, v2)
@@ -111,8 +124,8 @@ impl Metric for DotProductMetric {
fn similarity(&self, v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if is_x86_feature_detected!("avx2") == 0 {
return dot_similarity_avx2(v1, v2);
if is_x86_feature_detected!("avx2") {
return unsafe { dot_similarity_avx2(v1, v2) };
}
}
dot_similarity(v1, v2)
@@ -131,8 +144,8 @@ impl Metric for CosineMetric {
fn similarity(&self, v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if is_x86_feature_detected!("avx2") == 0 {
return dot_similarity_avx2(v1, v2);
if is_x86_feature_detected!("avx2") {
return unsafe { dot_similarity_avx2(v1, v2) };
}
}
dot_similarity(v1, v2)
@@ -141,8 +154,8 @@ impl Metric for CosineMetric {
fn preprocess(&self, vector: &[VectorElementType]) -> Option<Vec<VectorElementType>> {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if is_x86_feature_detected!("avx2") == 0 {
return Some(cosine_preprocess_avx2(vector));
if is_x86_feature_detected!("avx2") {
return Some(unsafe { cosine_preprocess_avx2(vector) });
}
}
Some(cosine_preprocess(vector))
@@ -159,4 +172,18 @@ mod tests {
let res = metric.preprocess(&[0.0, 0.0, 0.0, 0.0]);
eprintln!("res = {:#?}", res);
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[test]
fn test_avx2() {
if is_x86_feature_detected!("avx2") {
let v1 : Vec<f32> = vec![10., 11., 12., 13., 14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27., 28., 29., 30.];
let v2 : Vec<f32> = vec![40., 41., 42., 43., 44., 45., 46., 47., 48., 49., 50., 51., 52., 53., 54., 55., 56., 57., 58., 59., 60.];
let res1 = unsafe { euclid_similarity_avx2(&v1, &v2) };
let res2 = euclid_similarity(&v1, &v2);
println!("AVX2 = {}, orig = {}", res1, res2);
} else {
println!("AVX2 test skiped");
}
}
}