fix: vectorize as much code as possible (#569)

Signed-off-by: usamoi <[email protected]>

Author: usamoi

Cargo.lock

@@ -23,10 +23,10 @@ arrayvec.workspace = true
 bincode.workspace = true
 byteorder.workspace = true
 chrono = "0.4.38"
+half.workspace = true
 libc.workspace = true
 log.workspace = true
 memmap2.workspace = true
-num-traits.workspace = true
 paste.workspace = true
 pgrx = { version = "=0.12.1", default-features = false, features = [] }
 rand.workspace = true
@@ -67,11 +67,10 @@ arc-swap = "1.7.0"
 arrayvec = "0.7.4"
 bincode = "1.3.3"
 byteorder = "1.5.0"
-half = { version = "2.4.0", features = ["num-traits", "rand_distr", "serde"] }
+half = { version = "2.4.0", features = ["rand_distr", "serde"] }
 libc = "0.2.153"
 log = { version = "0.4.21", features = ["std"] }
 memmap2 = "0.9.4"
-num-traits = "0.2.19"
 parking_lot = "0.12.1"
 paste = "1.0.14"
 rand = "0.8.5"
@@ -99,3 +98,4 @@ overflow-checks = false
 [profile.release]
 lto = "fat"
 codegen-units = 1
+debug = true

Cargo.toml

@@ -6,15 +6,13 @@ edition.workspace = true
 [dependencies]
 half.workspace = true
 libc.workspace = true
-num-traits.workspace = true
 rand.workspace = true
 serde.workspace = true
 thiserror.workspace = true
 toml.workspace = true
 validator.workspace = true
 
 base_macros = { path = "../base_macros" }
-ccc = { path = "../ccc" }
 detect = { path = "../detect" }
 
 [lints]

crates/base/Cargo.toml

@@ -0,0 +1,3 @@
+#[derive(Debug, Clone, Copy)]
+#[repr(C, align(32))]
+pub struct Aligned32<T>(pub T);

crates/base/src/aligned.rs

@@ -1,13 +1,13 @@
-#![feature(core_intrinsics)]
 #![feature(avx512_target_feature)]
 #![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))]
-#![allow(internal_features)]
+#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512_f16))]
 #![allow(clippy::derivable_impls)]
 #![allow(clippy::len_without_is_empty)]
 #![allow(clippy::len_zero)]
 #![allow(clippy::needless_range_loop)]
 #![allow(clippy::nonminimal_bool)]
 
+pub mod aligned;
 pub mod always_equal;
 pub mod distance;
 pub mod index;

crates/base/src/lib.rs

@@ -6,9 +6,7 @@ use crate::vector::*;
 pub enum BVectorDot {}
 
 impl Operator for BVectorDot {
-    type VectorOwned = BVectorOwned;
-
-    const DISTANCE_KIND: DistanceKind = DistanceKind::Dot;
+    type Vector = BVectOwned;
 
     fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance {
         lhs.operator_dot(rhs)

crates/base/src/operator/bvector_dot.rs renamed to crates/base/src/operator/bvect_dot.rs

@@ -6,9 +6,7 @@ use crate::vector::*;
 pub enum BVectorHamming {}
 
 impl Operator for BVectorHamming {
-    type VectorOwned = BVectorOwned;
-
-    const DISTANCE_KIND: DistanceKind = DistanceKind::Hamming;
+    type Vector = BVectOwned;
 
     fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance {
         lhs.operator_hamming(rhs)

crates/base/src/operator/bvector_hamming.rs renamed to crates/base/src/operator/bvect_hamming.rs

@@ -6,9 +6,7 @@ use crate::vector::*;
 pub enum BVectorJaccard {}
 
 impl Operator for BVectorJaccard {
-    type VectorOwned = BVectorOwned;
-
-    const DISTANCE_KIND: DistanceKind = DistanceKind::Jaccard;
+    type Vector = BVectOwned;
 
     fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance {
         lhs.operator_jaccard(rhs)

crates/base/src/operator/bvector_jaccard.rs renamed to crates/base/src/operator/bvect_jaccard.rs

@@ -1,34 +1,27 @@
-mod bvector_dot;
-mod bvector_hamming;
-mod bvector_jaccard;
-mod svecf32_dot;
-mod svecf32_l2;
-mod vecf16_dot;
-mod vecf16_l2;
-mod vecf32_dot;
-mod vecf32_l2;
+mod bvect_dot;
+mod bvect_hamming;
+mod bvect_jaccard;
+mod svect_dot;
+mod svect_l2;
+mod vect_dot;
+mod vect_l2;
 
-pub use bvector_dot::BVectorDot;
-pub use bvector_hamming::BVectorHamming;
-pub use bvector_jaccard::BVectorJaccard;
-pub use svecf32_dot::SVecf32Dot;
-pub use svecf32_l2::SVecf32L2;
-pub use vecf16_dot::Vecf16Dot;
-pub use vecf16_l2::Vecf16L2;
-pub use vecf32_dot::Vecf32Dot;
-pub use vecf32_l2::Vecf32L2;
+pub use bvect_dot::BVectorDot;
+pub use bvect_hamming::BVectorHamming;
+pub use bvect_jaccard::BVectorJaccard;
+pub use svect_dot::SVectDot;
+pub use svect_l2::SVectL2;
+pub use vect_dot::VectDot;
+pub use vect_l2::VectL2;
 
 use crate::distance::*;
 use crate::vector::*;
 
 pub trait Operator: Copy + 'static + Send + Sync {
-    type VectorOwned: VectorOwned;
-
-    const DISTANCE_KIND: DistanceKind;
+    type Vector: VectorOwned;
 
     fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance;
 }
 
-pub type Owned<T> = <T as Operator>::VectorOwned;
-pub type Borrowed<'a, T> = <<T as Operator>::VectorOwned as VectorOwned>::Borrowed<'a>;
-pub type Scalar<T> = <<T as Operator>::VectorOwned as VectorOwned>::Scalar;
+pub type Owned<T> = <T as Operator>::Vector;
+pub type Borrowed<'a, T> = <<T as Operator>::Vector as VectorOwned>::Borrowed<'a>;

crates/base/src/operator/mod.rs

@@ -1,14 +1,14 @@
 use crate::distance::*;
 use crate::operator::*;
+use crate::scalar::ScalarLike;
 use crate::vector::*;
+use std::marker::PhantomData;
 
 #[derive(Debug, Clone, Copy)]
-pub enum SVecf32Dot {}
+pub struct SVectDot<S>(std::convert::Infallible, PhantomData<fn(S) -> S>);
 
-impl Operator for SVecf32Dot {
-    type VectorOwned = SVecf32Owned;
-
-    const DISTANCE_KIND: DistanceKind = DistanceKind::Dot;
+impl<S: ScalarLike> Operator for SVectDot<S> {
+    type Vector = SVectOwned<S>;
 
     fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance {
         lhs.operator_dot(rhs)

crates/base/src/operator/svecf32_l2.rs

@@ -0,0 +1,16 @@
+use crate::distance::*;
+use crate::operator::*;
+use crate::scalar::ScalarLike;
+use crate::vector::*;
+use std::marker::PhantomData;
+
+#[derive(Debug, Clone, Copy)]
+pub struct SVectL2<S>(std::convert::Infallible, PhantomData<fn(S) -> S>);
+
+impl<S: ScalarLike> Operator for SVectL2<S> {
+    type Vector = SVectOwned<S>;
+
+    fn distance(lhs: SVectBorrowed<'_, S>, rhs: SVectBorrowed<'_, S>) -> Distance {
+        lhs.operator_l2(rhs)
+    }
+}