Add reproducible K-means performance benchmark

diskann-benchmark/Cargo.toml

@@ -82,3 +82,6 @@ disk-index = [
     "dep:opentelemetry_sdk",
     "dep:scopeguard",
 ]
+
+# Compare the Disk and quantization K-means implementations.
+kmeans-comparison = ["dep:diskann-disk"]

diskann-benchmark/example/issue-939-kmeans-disk.json

@@ -0,0 +1,33 @@
+{
+  "search_directories": [],
+  "jobs": [
+    {
+      "type": "kmeans-comparison",
+      "content": {
+        "implementation": "disk",
+        "phase": "all",
+        "num_points": 50000,
+        "dimensions": [4, 32, 128, 384, 768, 1024, 3072],
+        "center_counts": [256],
+        "max_iterations": 10,
+        "thread_counts": [1],
+        "measurements": 30,
+        "seed": 42
+      }
+    },
+    {
+      "type": "kmeans-comparison",
+      "content": {
+        "implementation": "disk",
+        "phase": "all",
+        "num_points": 50000,
+        "dimensions": [4, 32, 128],
+        "center_counts": [256],
+        "max_iterations": 10,
+        "thread_counts": [2, 4, 8],
+        "measurements": 30,
+        "seed": 42
+      }
+    }
+  ]
+}

diskann-benchmark/example/issue-939-kmeans-init.json

@@ -0,0 +1,33 @@
+{
+  "search_directories": [],
+  "jobs": [
+    {
+      "type": "kmeans-comparison",
+      "content": {
+        "implementation": "disk",
+        "phase": "init",
+        "num_points": 50000,
+        "dimensions": [4, 32, 128, 384, 768, 1024, 3072],
+        "center_counts": [256],
+        "max_iterations": 10,
+        "thread_counts": [1],
+        "measurements": 30,
+        "seed": 42
+      }
+    },
+    {
+      "type": "kmeans-comparison",
+      "content": {
+        "implementation": "quantization",
+        "phase": "init",
+        "num_points": 50000,
+        "dimensions": [4, 32, 128, 384, 768, 1024, 3072],
+        "center_counts": [256],
+        "max_iterations": 10,
+        "thread_counts": [1],
+        "measurements": 30,
+        "seed": 42
+      }
+    }
+  ]
+}

diskann-benchmark/example/issue-939-kmeans-quantization.json

@@ -0,0 +1,33 @@
+{
+  "search_directories": [],
+  "jobs": [
+    {
+      "type": "kmeans-comparison",
+      "content": {
+        "implementation": "quantization",
+        "phase": "all",
+        "num_points": 50000,
+        "dimensions": [4, 32, 128, 384, 768, 1024, 3072],
+        "center_counts": [256],
+        "max_iterations": 10,
+        "thread_counts": [1],
+        "measurements": 30,
+        "seed": 42
+      }
+    },
+    {
+      "type": "kmeans-comparison",
+      "content": {
+        "implementation": "quantization",
+        "phase": "all",
+        "num_points": 50000,
+        "dimensions": [4, 32, 128],
+        "center_counts": [256],
+        "max_iterations": 10,
+        "thread_counts": [2, 4, 8],
+        "measurements": 30,
+        "seed": 42
+      }
+    }
+  ]
+}

diskann-benchmark/src/inputs/kmeans.rs

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::num::NonZeroUsize;
+
+use diskann_benchmark_runner::Checker;
+use serde::{Deserialize, Serialize};
+
+use super::{as_input, Example};
+
+#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub(crate) enum KmeansImplementation {
+    Disk,
+    Quantization,
+}
+
+impl std::fmt::Display for KmeansImplementation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Disk => write!(f, "disk"),
+            Self::Quantization => write!(f, "quantization"),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub(crate) enum KmeansPhase {
+    All,
+    Init,
+}
+
+impl std::fmt::Display for KmeansPhase {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::All => write!(f, "all"),
+            Self::Init => write!(f, "init"),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub(crate) struct KmeansComparison {
+    pub(crate) implementation: KmeansImplementation,
+    pub(crate) phase: KmeansPhase,
+    pub(crate) num_points: NonZeroUsize,
+    pub(crate) dimensions: Vec<NonZeroUsize>,
+    pub(crate) center_counts: Vec<NonZeroUsize>,
+    pub(crate) max_iterations: NonZeroUsize,
+    pub(crate) thread_counts: Vec<NonZeroUsize>,
+    pub(crate) measurements: NonZeroUsize,
+    pub(crate) seed: u64,
+}
+
+impl KmeansComparison {
+    pub(crate) const fn tag() -> &'static str {
+        "kmeans-comparison"
+    }
+
+    pub(crate) fn validate(&mut self, _checker: &mut Checker) -> anyhow::Result<()> {
+        anyhow::ensure!(!self.dimensions.is_empty(), "dimensions cannot be empty");
+        anyhow::ensure!(
+            !self.center_counts.is_empty(),
+            "center_counts cannot be empty"
+        );
+        anyhow::ensure!(
+            !self.thread_counts.is_empty(),
+            "thread_counts cannot be empty"
+        );
+        anyhow::ensure!(
+            self.center_counts
+                .iter()
+                .all(|count| count.get() <= self.num_points.get()),
+            "center counts cannot exceed num_points"
+        );
+        Ok(())
+    }
+}
+
+impl Example for KmeansComparison {
+    fn example() -> Self {
+        Self {
+            implementation: KmeansImplementation::Quantization,
+            phase: KmeansPhase::All,
+            num_points: NonZeroUsize::new(10_000).unwrap(),
+            dimensions: vec![NonZeroUsize::new(128).unwrap()],
+            center_counts: vec![NonZeroUsize::new(64).unwrap()],
+            max_iterations: NonZeroUsize::new(3).unwrap(),
+            thread_counts: vec![NonZeroUsize::new(1).unwrap()],
+            measurements: NonZeroUsize::new(10).unwrap(),
+            seed: 42,
+        }
+    }
+}
+
+as_input!(KmeansComparison);
+
+impl std::fmt::Display for KmeansComparison {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "K-means Implementation Benchmark\n")?;
+        writeln!(f, "{:>18}: {}", "implementation", self.implementation)?;
+        writeln!(f, "{:>18}: {}", "phase", self.phase)?;
+        writeln!(f, "{:>18}: {}", "points", self.num_points)?;
+        writeln!(f, "{:>18}: {:?}", "dimensions", self.dimensions)?;
+        writeln!(f, "{:>18}: {:?}", "center counts", self.center_counts)?;
+        writeln!(f, "{:>18}: {}", "max iterations", self.max_iterations)?;
+        writeln!(f, "{:>18}: {:?}", "thread counts", self.thread_counts)?;
+        writeln!(f, "{:>18}: {}", "measurements", self.measurements)?;
+        writeln!(f, "{:>18}: {}", "seed", self.seed)
+    }
+}

diskann-benchmark/src/inputs/mod.rs

@@ -8,6 +8,7 @@ pub(crate) mod exhaustive;
 pub(crate) mod filters;
 pub(crate) mod flat;
 pub(crate) mod graph_index;
+pub(crate) mod kmeans;
 pub(crate) mod multi_vector;
 pub(crate) mod post_processor;
 pub(crate) mod save_and_load;