perf(kernel): avoid per-element allocation and exploit kernel symmetry

algorithms/linfa-kernel/Cargo.toml

@@ -30,3 +30,13 @@ sprs = { version = "=0.11.2", default-features = false }
 
 linfa = { version = "0.8.1", path = "../.." }
 linfa-nn = { version = "0.8.1", path = "../linfa-nn" }
+
+[dev-dependencies]
+criterion = "0.5"
+ndarray-rand = "0.15"
+rand_xoshiro = "0.6"
+linfa = { version = "0.8.1", path = "../..", features = ["benchmarks"] }
+
+[[bench]]
+name = "kernel"
+harness = false

algorithms/linfa-kernel/benches/kernel.rs

@@ -0,0 +1,66 @@
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use linfa::benchmarks::config;
+use linfa::traits::Transformer;
+use linfa_kernel::{Kernel, KernelMethod, KernelType};
+use ndarray::Array2;
+use ndarray_rand::{rand::SeedableRng, rand_distr::Uniform, RandomExt};
+use rand_xoshiro::Xoshiro256Plus;
+
+/// Dense kernel construction is quadratic in both time and memory, so the
+/// sample sizes stay below the workspace-wide defaults: a 20_000 sample dense
+/// kernel would already need ~3.2 GB just to hold the matrix.
+const SIZES: [usize; 3] = [1_000, 2_000, 4_000];
+const N_FEATURES: [usize; 2] = [3, 8];
+
+fn dataset(n_samples: usize, n_features: usize) -> Array2<f64> {
+    let mut rng = Xoshiro256Plus::seed_from_u64(42);
+    Array2::random_using((n_samples, n_features), Uniform::new(-1.0, 1.0), &mut rng)
+}
+
+fn build_dense(dataset: &Array2<f64>, method: KernelMethod<f64>) {
+    Kernel::params()
+        .kind(KernelType::Dense)
+        .method(method)
+        .transform(dataset.view());
+}
+
+fn bench(c: &mut Criterion) {
+    let methods = [
+        (KernelMethod::Gaussian(0.5), "Gaussian"),
+        (KernelMethod::Linear, "Linear"),
+        (KernelMethod::Polynomial(1.0, 3.0), "Polynomial"),
+    ];
+
+    let mut group = c.benchmark_group("Kernel");
+    config::set_default_benchmark_configs(&mut group);
+
+    for (method, name) in methods {
+        for n_features in N_FEATURES {
+            for n_samples in SIZES {
+                let data = dataset(n_samples, n_features);
+
+                group.bench_with_input(
+                    BenchmarkId::new(
+                        format!("Dense-{name}-{n_features}feats"),
+                        format!("{n_samples}samples"),
+                    ),
+                    &data,
+                    |b, data| b.iter(|| build_dense(data, method.clone())),
+                );
+            }
+        }
+    }
+
+    group.finish();
+}
+
+#[cfg(not(target_os = "windows"))]
+criterion_group! {
+    name = benches;
+    config = config::get_default_profiling_configs();
+    targets = bench
+}
+#[cfg(target_os = "windows")]
+criterion_group!(benches, bench);
+
+criterion_main!(benches);

algorithms/linfa-kernel/src/lib.rs

@@ -25,7 +25,6 @@ use ndarray::Data;
 #[cfg(feature = "serde")]
 use serde_crate::{Deserialize, Serialize};
 use sprs::{CsMat, CsMatView};
-use std::ops::Mul;
 
 use linfa::{
     dataset::AsTargets, dataset::DatasetBase, dataset::FromTargetArray, dataset::Records,
@@ -281,8 +280,8 @@ impl<F: Float> KernelMethod<F> {
 
                 (-distance / eps).exp()
             }
-            KernelMethod::Linear => a.mul(&b).sum(),
-            KernelMethod::Polynomial(c, d) => (a.mul(&b).sum() + c).powf(d),
+            KernelMethod::Linear => a.dot(&b),
+            KernelMethod::Polynomial(c, d) => (a.dot(&b) + c).powf(d),
         }
     }
 
@@ -482,14 +481,19 @@ fn dense_from_fn<F: Float, D: Data<Elem = F>>(
     method: &KernelMethod<F>,
 ) -> Array2<F> {
     let n_observations = dataset.len_of(Axis(0));
-    let mut similarity = Array2::eye(n_observations);
+    let mut similarity = Array2::zeros((n_observations, n_observations));
 
+    // Every `KernelMethod` is symmetric in its two arguments, so only the upper
+    // triangle has to be evaluated and the result mirrored into the lower one.
     for i in 0..n_observations {
-        for j in 0..n_observations {
-            let a = dataset.row(i);
+        let a = dataset.row(i);
+
+        for j in i..n_observations {
             let b = dataset.row(j);
 
-            similarity[(i, j)] = method.distance(a, b);
+            let similarity_ij = method.distance(a, b);
+            similarity[(i, j)] = similarity_ij;
+            similarity[(j, i)] = similarity_ij;
         }
     }
 
@@ -530,6 +534,7 @@ mod tests {
     use linfa_nn::{BallTree, KdTree};
     use ndarray::{Array1, Array2};
     use std::f64::consts;
+    use std::ops::Mul;
 
     #[test]
     fn autotraits() {
@@ -614,6 +619,38 @@ mod tests {
         }
     }
 
+    #[test]
+    fn dense_from_fn_is_symmetric() {
+        // `dense_from_fn` only evaluates the upper triangle and mirrors it, which
+        // is valid because every `KernelMethod` is symmetric in its arguments.
+        // Guard both halves of that property.
+        let input_arr =
+            Array2::from_shape_fn((16, 4), |(i, j)| ((i * 31 + j * 17) % 100) as f64 / 100.);
+
+        for method in [
+            KernelMethod::Gaussian(0.5),
+            KernelMethod::Linear,
+            KernelMethod::Polynomial(1., 3.),
+        ] {
+            // the method itself is symmetric
+            for i in 0..16 {
+                for j in 0..16 {
+                    let forward = method.distance(input_arr.row(i), input_arr.row(j));
+                    let backward = method.distance(input_arr.row(j), input_arr.row(i));
+                    assert!((forward - backward).abs() <= 1e-12);
+                }
+            }
+
+            // and so is the matrix built from it
+            let similarity_matrix = dense_from_fn(&input_arr, &method);
+            for i in 0..16 {
+                for j in 0..16 {
+                    assert!((similarity_matrix[(i, j)] - similarity_matrix[(j, i)]).abs() <= 1e-12);
+                }
+            }
+        }
+    }
+
     #[test]
     fn gaussian_test() {
         let gauss_1 = KernelMethod::Gaussian(1.);