From ad7fd8fca20d7ffdf9a86de3302288e963bbe6e9 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 30 Jun 2026 11:43:07 +0200 Subject: [PATCH 01/38] feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: embedding clustering feat: embedding clustering feat: embedding clustering feat: embedding clustering feat: checkpoint feat: checkpoint feat: checkpoint fix: merge feat: checkpoint feat: checkpoint feat: checkpoint fix: merge feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint feat: checkpoint] feat: checkpoint] feat: checkpoint] feat: checkpoint feat: checkpoint --- libs/@local/graph/api/src/rest/entity/mod.rs | 63 +- .../store/postgres/knowledge/entity/mod.rs | 191 ++- .../graph/store/src/embedding/clustering.rs | 1437 +++++++++++++++++ .../graph/store/src/embedding/dimension.rs | 79 + .../graph/store/src/embedding/kernel.rs | 769 +++++++++ libs/@local/graph/store/src/embedding/mod.rs | 15 + libs/@local/graph/store/src/entity/mod.rs | 16 +- libs/@local/graph/store/src/entity/store.rs | 74 +- libs/@local/graph/store/src/error.rs | 15 + libs/@local/graph/store/src/lib.rs | 5 + libs/@local/graph/type-fetcher/src/store.rs | 20 +- tests/graph/integration/postgres/lib.rs | 11 + 12 files changed, 2663 insertions(+), 32 deletions(-) create mode 100644 libs/@local/graph/store/src/embedding/clustering.rs create mode 100644 libs/@local/graph/store/src/embedding/dimension.rs create mode 100644 libs/@local/graph/store/src/embedding/kernel.rs create mode 100644 libs/@local/graph/store/src/embedding/mod.rs diff --git a/libs/@local/graph/api/src/rest/entity/mod.rs b/libs/@local/graph/api/src/rest/entity/mod.rs index 1be3e54b8de..afc026cf3f1 100644 --- a/libs/@local/graph/api/src/rest/entity/mod.rs +++ b/libs/@local/graph/api/src/rest/entity/mod.rs @@ -13,14 +13,14 @@ use hash_graph_postgres_store::store::error::{EntityDoesNotExist, RaceConditionO use hash_graph_store::{ self, entity::{ - ClosedMultiEntityTypeMap, CreateEntityParams, DiffEntityParams, DiffEntityResult, - EntityPermissions, EntityQueryCursor, EntityQuerySortingRecord, EntityQuerySortingToken, - EntityQueryToken, EntityStore, EntityTypesError, EntityValidationReport, - EntityValidationType, HasPermissionForEntitiesParams, LinkDataStateError, - LinkDataValidationReport, LinkError, LinkTargetError, LinkValidationReport, - LinkedEntityError, MetadataValidationReport, PatchEntityParams, - PropertyMetadataValidationReport, QueryConversion, QueryEntitiesResponse, - SearchEntitiesFilter, SearchEntitiesParams, SearchEntitiesResponse, + ClosedMultiEntityTypeMap, ClusterEntitiesParams, ClusterEntitiesResponse, + CreateEntityParams, DiffEntityParams, DiffEntityResult, EntityCluster, EntityPermissions, + EntityQueryCursor, EntityQuerySortingRecord, EntityQuerySortingToken, EntityQueryToken, + EntityStore, EntityTypesError, EntityValidationReport, EntityValidationType, + HasPermissionForEntitiesParams, LinkDataStateError, LinkDataValidationReport, LinkError, + LinkTargetError, LinkValidationReport, LinkedEntityError, MetadataValidationReport, + PatchEntityParams, PropertyMetadataValidationReport, QueryConversion, + QueryEntitiesResponse, SearchEntitiesFilter, SearchEntitiesParams, SearchEntitiesResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, UnexpectedEntityType, UpdateEntityEmbeddingsParams, ValidateEntityComponents, ValidateEntityParams, }, @@ -97,6 +97,7 @@ use crate::rest::{ search_entities, patch_entity, update_entity_embeddings, + cluster_entities, diff_entity, ), components( @@ -114,6 +115,9 @@ use crate::rest::{ Embedding, UpdateEntityEmbeddingsParams, EntityEmbedding, + ClusterEntitiesParams, + ClusterEntitiesResponse, + EntityCluster, EntityQueryToken, PatchEntityParams, @@ -226,7 +230,12 @@ impl EntityResource { .route("/bulk", post(create_entities::)) .route("/diff", post(diff_entity::)) .route("/validate", post(validate_entity::)) - .route("/embeddings", post(update_entity_embeddings::)) + .nest( + "/embeddings", + Router::new() + .route("/", post(update_entity_embeddings::)) + .route("/clusters", post(cluster_entities::)), + ) .route("/permissions", post(has_permission_for_entities::)) .route("/search", post(search_entities::)) .nest( @@ -597,6 +606,42 @@ where .map_err(report_to_response) } +#[utoipa::path( + post, + path = "/entities/embeddings/clusters", + tag = "Entity", + params( + ("X-Authenticated-User-Actor-Id" = ActorEntityUuid, Header, description = "The ID of the actor which is used to authorize the request"), + ), + responses( + (status = 200, content_type = "application/json", description = "Clusters of entities by embedding similarity", body = ClusterEntitiesResponse), + (status = 422, content_type = "text/plain", description = "Provided request body is invalid"), + + (status = 500, description = "Store error occurred"), + ), + request_body = ClusterEntitiesParams, +)] +async fn cluster_entities( + AuthenticatedUserHeader(actor_id): AuthenticatedUserHeader, + store_pool: Extension>, + temporal_client: Extension>>, + Json(params): Json, +) -> Result, BoxedResponse> +where + S: StorePool + Send + Sync, +{ + let store = store_pool + .acquire(temporal_client.0) + .await + .map_err(report_to_response)?; + + store + .cluster_entities(actor_id, params) + .await + .map_err(report_to_response) + .map(Json) +} + #[utoipa::path( post, path = "/entities/diff", diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 3a41b84b8c7..ff70b0a7d00 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -2,6 +2,7 @@ mod delete; mod query; mod read; mod summary; + use alloc::borrow::Cow; use core::{borrow::Borrow as _, mem}; use std::collections::{HashMap, HashSet}; @@ -17,18 +18,21 @@ use hash_graph_authorization::policies::{ store::{PolicyCreationParams, PrincipalStore as _}, }; use hash_graph_store::{ + embedding::dimension::Dimension, entity::{ - CreateEntityParams, DeleteEntitiesParams, DeletionSummary, EmptyEntityTypes, - EntityPermissions, EntityQueryCursor, EntityQueryPath, EntityQuerySorting, EntityStore, - EntityTypeRetrieval, EntityTypesError, EntityValidationReport, EntityValidationType, - HasPermissionForEntitiesParams, PatchEntityParams, QueryConversion, QueryEntitiesParams, - QueryEntitiesResponse, QueryEntitySubgraphParams, QueryEntitySubgraphResponse, - SearchEntitiesFilter, SearchEntitiesParams, SearchEntitiesResponse, - SummarizeEntitiesParams, SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, - ValidateEntityComponents, ValidateEntityParams, + ClusterEntitiesParams, ClusterEntitiesResponse, CreateEntityParams, DeleteEntitiesParams, + DeletionSummary, EmptyEntityTypes, EntityCluster, EntityPermissions, EntityQueryCursor, + EntityQueryPath, EntityQuerySorting, EntityStore, EntityTypeRetrieval, EntityTypesError, + EntityValidationReport, EntityValidationType, HasPermissionForEntitiesParams, + PatchEntityParams, QueryConversion, QueryEntitiesParams, QueryEntitiesResponse, + QueryEntitySubgraphParams, QueryEntitySubgraphResponse, SummarizeEntitiesParams, + SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, ValidateEntityComponents, + ValidateEntityParams, }, entity_type::{EntityTypeStore as _, IncludeEntityTypeOption}, - error::{CheckPermissionError, DeletionError, InsertionError, QueryError, UpdateError}, + error::{ + CheckPermissionError, ClusterError, DeletionError, InsertionError, QueryError, UpdateError, + }, filter::{ Filter, FilterExpression, FilterExpressionList, Parameter, ParameterList, protection::transform_filter, @@ -2593,6 +2597,175 @@ where Ok(permitted_ids) } + + #[expect(clippy::too_many_lines)] + #[tracing::instrument(skip(self, params))] + async fn cluster_entities( + &self, + actor_id: ActorEntityUuid, + params: ClusterEntitiesParams, + ) -> Result> { + // 3072 fits in u16; compile-time verified. + const { + assert!(Embedding::DIM <= u16::MAX as usize); + } + #[expect( + clippy::cast_possible_truncation, + reason = "guarded by the const assertion above" + )] + const STORED_DIM: u16 = Embedding::DIM as u16; + + let dim = Dimension::new(params.dimension).ok_or_else(|| { + Report::new(ClusterError::InvalidDimension { + dimension: params.dimension, + }) + })?; + + if dim.get() > STORED_DIM { + return Err(Report::new(ClusterError::DimensionTooLarge { + dimension: dim.get(), + max: STORED_DIM, + })); + } + let truncated_dim = usize::from(dim.get()); + + // Filter to entities the actor is allowed to view. + let permitted = self + .has_permission_for_entities( + AuthenticatedActor::from(actor_id), + HasPermissionForEntitiesParams { + action: ActionName::ViewEntity, + entity_ids: Cow::Borrowed(¶ms.entity_ids), + temporal_axes: QueryTemporalAxesUnresolved::TransactionTime { + pinned: PinnedTemporalAxisUnresolved::new(None), + variable: VariableTemporalAxisUnresolved::new(None, None), + }, + include_drafts: false, + }, + ) + .await + .change_context(ClusterError::Store)?; + + let permitted_ids: Vec = params + .entity_ids + .iter() + .filter(|&id| permitted.contains_key(id)) + .copied() + .collect(); + + let entity_uuids: Vec = permitted_ids.iter().map(|id| id.entity_uuid).collect(); + let web_ids: Vec = permitted_ids.iter().map(|id| id.web_id).collect(); + + // Truncate server-side via `subvector` so postgres only sends + // `truncated_dim`-dimensional vectors over the wire. + let row_stream = self + .as_client() + .query_raw( + &format!( + "SELECT + e.web_id, + e.entity_uuid, + subvector(e.embedding, 1, {truncated_dim})::vector({truncated_dim}) AS \ + embedding + FROM entity_embeddings e + WHERE e.property IS NULL + AND (e.web_id, e.entity_uuid) IN (SELECT unnest($1::uuid[]), \ + unnest($2::uuid[]))" + ), + [ + &web_ids as &(dyn ToSql + Sync), + &entity_uuids as &(dyn ToSql + Sync), + ], + ) + .instrument(tracing::info_span!( + "cluster_entities.embeddings", + otel.kind = "client", + db.system = "postgresql", + peer.service = "Postgres", + )) + .await + .change_context(ClusterError::Store)?; + + let mut row_stream = core::pin::pin!(row_stream); + + let mut flat: Vec = Vec::with_capacity(permitted_ids.len() * truncated_dim); + let mut found_ids: Vec = Vec::with_capacity(permitted_ids.len()); + + while let Some(row) = row_stream + .try_next() + .await + .change_context(ClusterError::Store)? + { + let web_id: WebId = row.get(0); + let entity_uuid: EntityUuid = row.get(1); + let embedding: Embedding<'_> = row.get(2); + + flat.extend(embedding.iter()); + + found_ids.push(EntityId { + web_id, + entity_uuid, + draft_id: None, + }); + } + + // Every requested entity not in a cluster goes into + // `missing_embeddings`, whether due to permissions or no embedding. + // Distinguishing the two would leak permission information. + let found_set: HashSet<(WebId, EntityUuid)> = found_ids + .iter() + .map(|id| (id.web_id, id.entity_uuid)) + .collect(); + let missing_embeddings: Vec = params + .entity_ids + .into_iter() + .filter(|id| !found_set.contains(&(id.web_id, id.entity_uuid))) + .collect(); + + if found_ids.is_empty() || params.cluster_count == 0 { + return Ok(ClusterEntitiesResponse { + clusters: Vec::new(), + missing_embeddings, + }); + } + + let config = hash_graph_store::embedding::clustering::Config::for_k_with_seed( + params.cluster_count, + params.seed.unwrap_or_else(|| { + std::time::SystemTime::UNIX_EPOCH + .elapsed() + .map_or(0, |elapsed| { + #[expect( + clippy::cast_possible_truncation, + reason = "seed only needs entropy, truncation is fine" + )] + let seed = elapsed.as_nanos() as u64; + seed + }) + }), + ); + + let result = hash_graph_store::embedding::clustering::cluster(&flat, dim, &config); + + let mut groups: HashMap> = HashMap::new(); + for (index, id) in found_ids.iter().enumerate() { + groups.entry(result.label(index)).or_default().push(*id); + } + + let clusters = groups + .into_iter() + .map(|(cluster_id, entity_ids)| EntityCluster { + cluster_id, + entity_ids, + centroid: result.centroid(cluster_id).to_vec(), + }) + .collect(); + + Ok(ClusterEntitiesResponse { + clusters, + missing_embeddings, + }) + } } #[derive(Debug)] diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs new file mode 100644 index 00000000000..f48da06398a --- /dev/null +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -0,0 +1,1437 @@ +use alloc::borrow::Cow; +use core::{cmp, mem, num::NonZero}; + +use rand::{Rng, RngExt as _, SeedableRng as _}; +use rand_xoshiro::Xoshiro256PlusPlus; +use rayon::prelude::*; + +use super::{dimension::Dimension, kernel}; + +/// Parameters for k-means clustering. +/// +/// Use [`Config::for_k`] or [`Config::for_k_with_seed`] to construct with +/// reasonable defaults, then override individual fields as needed. +pub struct Config { + /// Number of clusters. + pub k: u16, + + /// Maximum Lloyd iterations per run before declaring convergence. + pub max_iters: NonZero, + + /// Number of independent restarts. The run with the lowest inertia wins. + pub n_init: NonZero, + + /// Convergence tolerance: a run stops early when the relative change in + /// inertia between iterations falls below this value. + pub tol: f32, + + /// Maximum number of points sampled during k-means++ seeding. + /// Capped to avoid quadratic seeding cost on very large datasets. + pub sample_cap: usize, + + /// Base seed for the PRNG. Each restart derives its own seed from this. + pub seed: u64, + + /// Number of points processed per batch in the assignment step. + pub chunk: NonZero, +} + +impl Config { + /// Creates a configuration for `k` clusters, drawing the seed from `rng`. + #[must_use] + pub(crate) fn for_k(k: u16, mut rng: impl Rng) -> Self { + Self::for_k_with_seed(k, rng.random()) + } + + /// Creates a configuration for `k` clusters with a fixed seed. + /// + /// Defaults: 30 max iterations, 5 restarts, 1e-4 convergence tolerance, + /// sample cap of min(256k, 8192), chunk size 256. + #[must_use] + pub fn for_k_with_seed(k: u16, seed: u64) -> Self { + Self { + k, + max_iters: const { NonZero::new(30).unwrap() }, + n_init: const { NonZero::new(5).unwrap() }, + tol: 1e-4, + sample_cap: cmp::min(256 * usize::from(k), 8192), + seed, + chunk: const { NonZero::new(256).unwrap() }, + } + } +} + +/// Result of spherical k-means clustering. +/// +/// `centroids` is a flat `k * d` row-major buffer where `d` is the +/// embedding [`Dimension`]. Centroid `i` occupies +/// `centroids[i * d .. (i + 1) * d]`. +pub struct Clustering { + pub dimension: Dimension, + + /// Flat centroid matrix, `k * d` elements in row-major order. + pub centroids: Box<[f32]>, + + /// Cluster assignment for each input point, values in `0..k`. + pub labels: Box<[u16]>, +} + +impl Clustering { + /// Allocates a zeroed clustering for `k` centroids over `n` points. + fn new(k: u16, n: usize, d: Dimension) -> Self { + // SAFETY: All-zero bits are valid for `f32` (IEEE 754 positive zero) + // and for `u16` (the integer 0). `Box::new_zeroed_slice` allocates + // zeroed memory of the correct layout, so `assume_init` is sound. + let centroids: Box<[f32]> = + unsafe { Box::new_zeroed_slice((k as usize) * (d.get() as usize)).assume_init() }; + // SAFETY: All-zero bits are valid for `u16` (the integer 0). + let labels: Box<[u16]> = unsafe { Box::new_zeroed_slice(n).assume_init() }; + + Self { + centroids, + labels, + dimension: d, + } + } + + /// Returns the `D`-dimensional slice for centroid `cluster`. + #[must_use] + pub fn centroid(&self, cluster: u16) -> &[f32] { + &self.centroids[cluster as usize * (self.dimension.get() as usize) + ..(cluster + 1) as usize * (self.dimension.get() as usize)] + } + + /// Returns a mutable `D`-dimensional slice for centroid `cluster`. + fn centroid_mut(&mut self, cluster: u16) -> &mut [f32] { + &mut self.centroids[cluster as usize * (self.dimension.get() as usize) + ..(cluster + 1) as usize * (self.dimension.get() as usize)] + } + + /// Returns the cluster label for point `entity`. + #[must_use] + pub fn label(&self, entity: usize) -> u16 { + self.labels[entity] + } + + /// Returns a mutable reference to the cluster label for point `entity`. + fn label_mut(&mut self, entity: usize) -> &mut u16 { + &mut self.labels[entity] + } +} + +// TODO: I wonder if we can make this allocation less +fn sample_indices(n: usize, m: usize, mut rng: impl Rng) -> Vec { + let mut idx: Vec = (0..n).collect(); + + for i in 0..m { + let j = i + rng.random_range(0..n - i); // partial Fisher–Yates + idx.swap(i, j); + } + + idx.truncate(m); + idx +} + +/// Squared chord distance between a point and a unit centroid. +/// +/// For a unit centroid `c` and a point with inverse norm `inv`, the cosine +/// similarity is `dot(point, c) * inv`. The squared chord distance is +/// `2 - 2 * similarity`, which lies in `[0, 4]` and equals `||u - c||²` +/// when `u` is the unit-normalized point. +/// +/// Returns `0.0` for zero-norm points (`point_inv_norm == 0.0`). +/// +/// This is a squared distance. Do not square it again for D² sampling. +#[inline] +fn squared_chord_distance(dot: f32, point_inv_norm: f32) -> f32 { + if point_inv_norm == 0.0 { + return 0.0; + } + + let similarity = (dot * point_inv_norm).clamp(-1.0, 1.0); + + 2.0_f32.mul_add(-similarity, 2.0).max(0.0) +} + +/// Finds the nearest centroid to `point` and returns its index and spherical +/// distance. +/// +/// # Safety +/// +/// * `point.len() == D` +/// * `centroids.len() == k * D` +/// * `k > 0` +/// * `D` is a multiple of 8 (enforced at compile time by the const generic). +#[inline] +#[must_use] +pub(crate) unsafe fn nearest_centroid( + point: &[f32], + point_inv_norm: f32, + centroids: &[f32], + k: usize, + d: usize, +) -> (u16, f32) { + debug_assert_eq!(point.len(), d); + debug_assert_eq!(centroids.len(), k * d); + debug_assert!(k > 0); + + // SAFETY: the caller guarantees these preconditions. The hints let the + // compiler elide bounds checks on the centroid slicing inside the loop. + unsafe { + core::hint::assert_unchecked(point.len() == d); + core::hint::assert_unchecked(centroids.len() == k * d); + core::hint::assert_unchecked(d.is_multiple_of(8)); + core::hint::assert_unchecked(k > 0); + } + + let mut best = 0; + let mut best_dot = f32::NEG_INFINITY; + + for cluster in 0..k { + let start = cluster * d; + let centroid = ¢roids[start..start + d]; + + // SAFETY: `point` and `centroid` both have length `D`, and `D` is a + // multiple of 8 (guaranteed by Dimension). + let dot = unsafe { kernel::dot(point, centroid) }; + + #[expect( + clippy::cast_possible_truncation, + reason = "k is supposed to be low, and checked as such via the config" + )] + if dot > best_dot { + best = cluster as u16; + best_dot = dot; + } + } + + (best, squared_chord_distance(best_dot, point_inv_norm)) +} + +/// Pre-allocated scratch space for the k-means fitting loop. +/// +/// All buffers are allocated once and reused across restarts to avoid +/// per-iteration allocation overhead. +struct Fit { + k: usize, + m: usize, + d: usize, + + /// Current centroids for this restart, `k * d` elements. + centroids: Box<[f32]>, + /// Best centroids seen across all restarts. + best_centroids: Box<[f32]>, + /// Per-cluster accumulator for centroid recomputation, `k * d` elements. + sums: Box<[f32]>, + /// Per-cluster point count for centroid averaging. + counts: Box<[usize]>, + /// Per-sample-point cluster assignment. + labels: Box<[u16]>, + /// Per-sample-point closest centroid distance (for k-means++ seeding). + closest_distances: Box<[f32]>, + /// Tracks which sample points have been selected as seeds. + selected: Box<[bool]>, + /// Lowest inertia across all restarts. + best_inertia: f32, +} + +impl Fit { + fn new(k: usize, m: usize, d: usize) -> Self { + // SAFETY: all-zero bits are valid for f32 (IEEE 754 +0.0), usize (0), u16 (0), and bool + // (false). `Box::new_zeroed_slice` allocates zeroed memory of the correct layout + // for each type, so `assume_init` is sound in every case. + let centroids = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; + // SAFETY: see above + let best_centroids = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; + // SAFETY: see above + let sums = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; + // SAFETY: see above + let counts = unsafe { Box::<[usize]>::new_zeroed_slice(k).assume_init() }; + // SAFETY: see above + let labels = unsafe { Box::<[u16]>::new_zeroed_slice(m).assume_init() }; + // SAFETY: see above + let closest_distances = unsafe { Box::<[f32]>::new_zeroed_slice(m).assume_init() }; + // SAFETY: see above + let selected = unsafe { Box::<[bool]>::new_zeroed_slice(m).assume_init() }; + let best_inertia = f32::INFINITY; + + Self { + k, + m, + d, + centroids, + best_centroids, + sums, + counts, + labels, + closest_distances, + selected, + best_inertia, + } + } + + fn reset_centroids(&mut self) { + self.centroids.fill(0.0); + } + + fn reset_sums(&mut self) { + self.sums.fill(0.0); + } + + fn reset_counts(&mut self) { + self.counts.fill(0); + } + + fn reset_selected(&mut self) { + self.selected.fill(false); + } + + /// Reinitializes empty clusters from the sample point farthest from + /// its assigned centroid. + /// + /// For each empty cluster, scans the sample to find the point with + /// the largest squared chord distance to its current centroid, copies + /// that point as the new centroid (normalized), and updates the + /// point's label so it won't be picked again for subsequent empty + /// clusters in the same pass. + #[expect( + clippy::cast_possible_truncation, + reason = "cluster index < k, and k originates from Config::k (u16)" + )] + fn reinit_empty_clusters(&mut self, sample: &[f32], sample_inv_norms: &[f32]) -> bool { + let &mut Self { d, k, .. } = self; + + let mut reseeded = false; + + for cluster in 0..k { + if self.counts[cluster] != 0 { + continue; + } + + reseeded = true; + let mut farthest_idx = 0; + let mut farthest_dist = -1.0_f32; + + for (i, (point, &inv_norm)) in sample.chunks_exact(d).zip(sample_inv_norms).enumerate() + { + let label = usize::from(self.labels[i]); + let c_start = label * d; + + // SAFETY: point and centroid both have length `d`, + // a multiple of 8 (guaranteed by Dimension). + let dot = unsafe { kernel::dot(point, &self.centroids[c_start..c_start + d]) }; + let dist = squared_chord_distance(dot, inv_norm); + + if dist > farthest_dist { + farthest_dist = dist; + farthest_idx = i; + } + } + + let point_start = farthest_idx * d; + let centroid_start = cluster * d; + self.centroids[centroid_start..centroid_start + d] + .copy_from_slice(&sample[point_start..point_start + d]); + + // SAFETY: centroid row has length `d`, a multiple of 8. + unsafe { + kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); + } + + // Update the label so the next empty cluster picks a different + // point (this point's distance to its new centroid is ~0). + self.labels[farthest_idx] = cluster as u16; + } + + reseeded + } + + /// Runs k-means++ initialization followed by Lloyd iterations on the + /// sample, repeating for `n_init` restarts. The best centroids (lowest + /// inertia) are stored in `self.best_centroids`. + fn run( + &mut self, + sample: &[f32], + chunk: usize, + row_chunk: usize, + sample_inv_norms: &[f32], + mut rng: impl Rng, + config: &Config, + ) { + for _ in 0..config.n_init.get() { + self.reset_centroids(); + self.closest_distances.fill(f32::INFINITY); + self.reset_selected(); + + self.seed_plusplus(sample, sample_inv_norms, &mut rng); + + let inertia = self.lloyd(sample, chunk, row_chunk, sample_inv_norms, config); + + if inertia < self.best_inertia { + self.best_inertia = inertia; + mem::swap(&mut self.best_centroids, &mut self.centroids); + } + } + } + + /// Runs Lloyd iterations on the sample until convergence or `max_iters`. + /// Returns the final inertia (sum of distances to assigned centroids). + fn lloyd( + &mut self, + sample: &[f32], + chunk: usize, + row_chunk: usize, + sample_inv_norms: &[f32], + config: &Config, + ) -> f32 { + let &mut Self { d, k, .. } = self; + let mut previous_inertia = f32::INFINITY; + let mut inertia = f32::INFINITY; + + for _ in 0..config.max_iters.get() { + inertia = sample + .par_chunks(row_chunk) + .zip(sample_inv_norms.par_chunks(chunk)) + .zip(self.labels.par_chunks_mut(chunk)) + .map(|((points, inv_norms), labels)| { + let mut inertia = 0.0; + let count = labels.len(); + + // SAFETY: each parallel chunk pairs `count` labels with + // `count * d` floats of point data and `count` inv_norms. + // `d` is a multiple of 8 (guaranteed by Dimension). + unsafe { + core::hint::assert_unchecked(points.len() == count * d); + core::hint::assert_unchecked(inv_norms.len() == count); + core::hint::assert_unchecked(d.is_multiple_of(8)); + } + + let mut i = 0; + while i + 4 <= count { + let p0 = &points[i * d..i * d + d]; + let p1 = &points[(i + 1) * d..(i + 1) * d + d]; + let p2 = &points[(i + 2) * d..(i + 2) * d + d]; + let p3 = &points[(i + 3) * d..(i + 3) * d + d]; + + // SAFETY: each point length d, centroids length k*d, + // k > 0, d a multiple of 8 (guaranteed by Dimension). + let nearest = + unsafe { kernel::nearest4(p0, p1, p2, p3, &self.centroids, k, d) }; + + let inv = [ + inv_norms[i], + inv_norms[i + 1], + inv_norms[i + 2], + inv_norms[i + 3], + ]; + for m in 0..4 { + labels[i + m] = nearest[m].0; + inertia += squared_chord_distance(nearest[m].1, inv[m]); + } + i += 4; + } + + while i < count { + let point = &points[i * d..i * d + d]; + // SAFETY: point length d, centroids length k*d, k > 0, + // d mult of 8. + let (label, distance) = + unsafe { nearest_centroid(point, inv_norms[i], &self.centroids, k, d) }; + labels[i] = label; + inertia += distance; + i += 1; + } + + inertia + }) + .sum(); + + self.reset_sums(); + self.reset_counts(); + + for ((point, label), inv_norm) in sample + .chunks_exact(d) + .zip(self.labels.iter().copied()) + .zip(sample_inv_norms.iter().copied()) + { + let cluster = usize::from(label); + let start = cluster * d; + + self.counts[cluster] += 1; + + if inv_norm == 0.0 { + continue; + } + + // SAFETY: `sums[start..start + d]` and `point` both have + // length `d`, and `d` is a multiple of 8 (guaranteed by Dimension). + unsafe { + kernel::add_scaled_into(&mut self.sums[start..start + d], point, inv_norm); + } + } + + for cluster in 0..k { + if self.counts[cluster] == 0 { + continue; + } + + let start = cluster * d; + let centroid = &mut self.centroids[start..start + d]; + let sum = &self.sums[start..start + d]; + + // SAFETY: `centroid` and `sum` both have length `D`, and `D` + // is a multiple of 8 (guaranteed by Dimension). + unsafe { + #[expect( + clippy::cast_precision_loss, + reason = "cluster count is bounded by sample_cap (≤8192), well within f32 \ + precision" + )] + let inv_count = 1.0 / self.counts[cluster] as f32; + kernel::scale_into(centroid, sum, inv_count); + } + + // SAFETY: centroid rows have length `D`, and `D` is a + // multiple of 8 (guaranteed by Dimension). + unsafe { + kernel::normalize(centroid); + } + } + + let reseeded = self.reinit_empty_clusters(sample, sample_inv_norms); + + // Skip the convergence check when a cluster was just reseeded: + // the reseeded centroid hasn't had an assignment pass yet, so + // breaking now would waste the reinit. + if !reseeded && previous_inertia.is_finite() { + let relative_change = + (previous_inertia - inertia).abs() / previous_inertia.max(f32::EPSILON); + + if relative_change <= config.tol { + break; + } + } + + previous_inertia = inertia; + } + + inertia + } + + /// k-means++ D² weighted seeding. Picks `k` initial centroids from the + /// sample, each chosen with probability proportional to its squared + /// distance from the nearest already-chosen centroid. + fn seed_plusplus(&mut self, sample: &[f32], sample_inv_norms: &[f32], mut rng: impl Rng) { + let &mut Self { d, k, m, .. } = self; + + let mut restart_rng = Xoshiro256PlusPlus::seed_from_u64(rng.random()); + let mut point = restart_rng.random_range(0..m); + + for cluster in 0..k { + let centroid_start = cluster * d; + let point_start = point * d; + + self.centroids[centroid_start..centroid_start + d] + .copy_from_slice(&sample[point_start..point_start + d]); + + // SAFETY: centroid rows have length `D`, and `D` is a multiple of 8 (guaranteed by + // Dimension). + unsafe { + kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); + } + + self.selected[point] = true; + + let centroid = &self.centroids[centroid_start..centroid_start + d]; + + let total: f32 = sample + .par_chunks_exact(d) + .zip(sample_inv_norms.par_iter().copied()) + .zip(self.closest_distances.par_iter_mut()) + .enumerate() + .map(|(index, ((point, inv_norm), closest))| { + if self.selected[index] { + *closest = 0.0; + return 0.0; + } + + // SAFETY: `point` and `centroid` both have length `D`, and + // `D` is a multiple of 8 (guaranteed by Dimension). + let dot = unsafe { kernel::dot(point, centroid) }; + let distance = squared_chord_distance(dot, inv_norm); + + if distance < *closest { + *closest = distance; + } + + *closest + }) + .sum(); + + if cluster + 1 == k { + break; + } + + point = if total.is_finite() && total > 0.0 { + let mut target = restart_rng.random_range(0.0..total); + let mut sampled = self + .closest_distances + .iter() + .rposition(|distance| *distance > 0.0) + .unwrap_or(0); + + for (index, distance) in self.closest_distances.iter().copied().enumerate() { + if distance <= 0.0 { + continue; + } + + target -= distance; + + if target <= 0.0 { + sampled = index; + break; + } + } + + sampled + } else { + let remaining = self.selected.iter().filter(|selected| !**selected).count(); + let mut target = restart_rng.random_range(0..remaining); + let mut sampled = 0; + + for (index, selected) in self.selected.iter().copied().enumerate() { + if selected { + continue; + } + + if target == 0 { + sampled = index; + break; + } + + target -= 1; + } + + sampled + }; + } + } +} + +/// Per-thread accumulator for parallel centroid recomputation. +/// +/// Each rayon task gets its own `Accum`; they are merged via [`Accum::merge`] +/// after the parallel fold completes. +struct Accum { + /// Per-cluster sum of normalized points, `k * d` elements. + sums: Box<[f32]>, + /// Per-cluster point count. + counts: Box<[usize]>, +} + +impl Accum { + fn new(k: usize, d: usize) -> Self { + // SAFETY: all-zero bits are valid for f32 (0.0) and usize (0). `assume_init` is + // sound after `new_zeroed_slice`. + let sums = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; + // SAFETY: see above + let counts = unsafe { Box::<[usize]>::new_zeroed_slice(k).assume_init() }; + Self { sums, counts } + } + + fn merge(mut self, other: &Self, k: usize, d: usize) -> Self { + for cluster in 0..k { + let start = cluster * d; + + self.counts[cluster] += other.counts[cluster]; + + // SAFETY: both cluster sum rows have length `d`, and `d` is a + // multiple of 8 (guaranteed by Dimension). + unsafe { + kernel::add_into( + &mut self.sums[start..start + d], + &other.sums[start..start + d], + ); + } + } + + self + } +} + +/// Assigns all `n` points to their nearest centroid, recomputes centroids +/// from the full population, and re-assigns labels to the final centroids. +/// +/// Uses a parallel fold/reduce: each rayon task accumulates into its own +/// [`Accum`], then results are merged. The final centroids are averaged +/// and normalized in-place. +/// +/// # Safety +/// +/// * `x.len() == n * D` for some `n` +/// * `clustering.centroids.len() == k * D` +/// * `clustering.labels.len() == n` +/// * `k > 0` +/// * `D` is a multiple of 8 (guaranteed by Dimension) +unsafe fn assign(x: &[f32], clustering: &mut Clustering, k: usize, chunk: usize, row_chunk: usize) { + let d = clustering.dimension.get() as usize; + + let full = x + .par_chunks(row_chunk) + .zip(clustering.labels.par_chunks_mut(chunk)) + .fold( + || Accum::new(k, d), + |mut accum, (points, labels)| { + // SAFETY: `cluster` established `x.len() == n * D` and + // `centroids.len() == k * D`. `par_chunks(row_chunk)` with + // `row_chunk = chunk * D` produces chunks where + // `points.len()` is a multiple of `D` and matches `labels.len() * D`. + unsafe { + assign_chunk(&clustering.centroids, k, d, points, labels, &mut accum); + } + + accum + }, + ) + .reduce(|| Accum::new(k, d), |lhs, rhs| lhs.merge(&rhs, k, d)); + + for cluster in 0..k { + if full.counts[cluster] == 0 { + continue; + } + + let start = cluster * d; + + #[expect( + clippy::cast_possible_truncation, + reason = "cluster < k and k originates from Config::k (u16)" + )] + let centroid = clustering.centroid_mut(cluster as u16); + let sum = &full.sums[start..start + d]; + + // SAFETY: centroid and sum both length D, a multiple of 8. + unsafe { + #[expect( + clippy::cast_precision_loss, + reason = "cluster count bounded by n; precision loss acceptable for averaging" + )] + let inv_count = 1.0 / full.counts[cluster] as f32; + kernel::scale_into(centroid, sum, inv_count); + } + // SAFETY: centroid length D, a multiple of 8. + unsafe { + kernel::normalize(centroid); + } + } + + // SAFETY: centroids were just recomputed; same invariants hold. + unsafe { + reassign( + x, + &clustering.centroids, + &mut clustering.labels, + k, + d, + chunk, + row_chunk, + ); + } +} + +/// Processes one parallel chunk of the assignment step: finds the nearest +/// centroid for each point, accumulates normalized points into cluster sums, +/// and records labels. +/// +/// # Safety +/// +/// * `points.len() == labels.len() * d` +/// * `centroids.len() >= k * d` +/// * `d` is a multiple of 8 +/// * `k > 0` +/// * `accum.sums.len() >= k * d` and `accum.counts.len() >= k` +unsafe fn assign_chunk( + centroids: &[f32], + k: usize, + d: usize, + points: &[f32], + labels: &mut [u16], + accum: &mut Accum, +) { + // field path -> disjoint capture of `centroids` only, leaving + // `labels` free for the mutable parallel borrow. + let count = labels.len(); + + // SAFETY: each parallel chunk pairs `count` labels with + // `count * D` floats of point data. `D` is a compile-time + // multiple of 8. + unsafe { + core::hint::assert_unchecked(points.len() == count * d); + core::hint::assert_unchecked(d.is_multiple_of(8)); + } + + let mut i = 0; + while i + 4 <= count { + let p0 = &points[i * d..i * d + d]; + let p1 = &points[(i + 1) * d..(i + 1) * d + d]; + let p2 = &points[(i + 2) * d..(i + 2) * d + d]; + let p3 = &points[(i + 3) * d..(i + 3) * d + d]; + + // SAFETY: each point length D, centroids length k*D, k > 0, D a multiple of 8 (guaranteed + // by Dimension). + let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; + let ps = [p0, p1, p2, p3]; + + for m in 0..4 { + let label = nearest[m].0; + labels[i + m] = label; + let cluster = usize::from(label); + accum.counts[cluster] += 1; + + let start = cluster * d; + + // SAFETY: point length D, a multiple of 8. + let norm = unsafe { kernel::dot(ps[m], ps[m]).sqrt() }; + if norm == 0.0 { + continue; + } + + // SAFETY: `sums[start..start + D]` and `point` both have length `D`, and `D` is a + // multiple of 8 (guaranteed by Dimension). + unsafe { + kernel::add_scaled_into(&mut accum.sums[start..start + d], ps[m], norm.recip()); + } + } + i += 4; + } + + while i < count { + let point = &points[i * d..i * d + d]; + // SAFETY: point length D, centroids length k*D, k > 0, D mult of 8. + let (label, _) = unsafe { nearest_centroid(point, 1.0, centroids, k, d) }; + labels[i] = label; + let cluster = usize::from(label); + accum.counts[cluster] += 1; + + let start = cluster * d; + + // SAFETY: point length D. + let norm = unsafe { kernel::dot(point, point).sqrt() }; + if norm != 0.0 { + // SAFETY: `sums[start..start + D]` and `point` both + // have length `D`, and `D` is a multiple of 8. + unsafe { + kernel::add_scaled_into(&mut accum.sums[start..start + d], point, norm.recip()); + } + } + i += 1; + } +} + +/// Processes one parallel chunk of the reassignment step: updates each +/// label to the nearest final centroid. +/// +/// # Safety +/// +/// * `points.len() == labels.len() * d` +/// * `centroids.len() >= k * d` +/// * `d` is a multiple of 8 +/// * `k > 0` +unsafe fn reassign_chunk( + k: usize, + d: usize, + centroids: &[f32], + points: &[f32], + labels: &mut [u16], +) { + let count = labels.len(); + + // SAFETY: each parallel chunk pairs `count` labels with + // `count * D` floats of point data. `D` is a compile-time + // multiple of 8. + unsafe { + core::hint::assert_unchecked(points.len() == count * d); + core::hint::assert_unchecked(d.is_multiple_of(8)); + } + + let mut i = 0; + while i + 4 <= count { + let p0 = &points[i * d..i * d + d]; + let p1 = &points[(i + 1) * d..(i + 1) * d + d]; + let p2 = &points[(i + 2) * d..(i + 2) * d + d]; + let p3 = &points[(i + 3) * d..(i + 3) * d + d]; + + // SAFETY: each point length D, centroids length k*D, k > 0, + // D a multiple of 8 (guaranteed by Dimension). + let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; + + labels[i] = nearest[0].0; + labels[i + 1] = nearest[1].0; + labels[i + 2] = nearest[2].0; + labels[i + 3] = nearest[3].0; + i += 4; + } + + while i < count { + let point = &points[i * d..i * d + d]; + // SAFETY: point length D, centroids length k*D, k > 0, D mult of 8. + let (label, _) = unsafe { nearest_centroid(point, 1.0, centroids, k, d) }; + labels[i] = label; + i += 1; + } +} + +/// Re-assigns labels to the nearest final centroid. +/// +/// After centroid recomputation, some boundary points may no longer be +/// nearest to the centroid stored under their label. This pass fixes that. +/// +/// # Safety +/// +/// Same as [`assign`]. +unsafe fn reassign( + x: &[f32], + centroids: &[f32], + labels: &mut [u16], + k: usize, + d: usize, + chunk: usize, + row_chunk: usize, +) { + x.par_chunks(row_chunk) + .zip(labels.par_chunks_mut(chunk)) + .for_each(|(points, labels)| { + // SAFETY: `par_chunks(row_chunk)` with `row_chunk = chunk * D` + // ensures `points.len() == labels.len() * D`. Centroids and k + // are valid from the caller. + unsafe { + reassign_chunk(k, d, centroids, points, labels); + } + }); +} + +/// Runs spherical k-means over a flat row-major embedding matrix. +/// +/// `x` contains `n` points of `dimension` floats each, laid out +/// contiguously. Returns cluster assignments and unit-normalized centroids. +/// +/// # Panics +/// +/// Panics if `x.len()` is not a multiple of `dimension`. +#[must_use] +#[expect(clippy::integer_division_remainder_used, clippy::integer_division)] +pub fn cluster(x: &[f32], dimension: Dimension, config: &Config) -> Clustering { + let d = dimension.get() as usize; + assert!(x.len().is_multiple_of(d)); + + let n = x.len() / d; + let k = cmp::min(config.k, n.saturating_truncate()); + + let mut clustering = Clustering::new(k, n, dimension); + + if k == 0 { + return clustering; + } + + let k = usize::from(k); + let mut rng = Xoshiro256PlusPlus::seed_from_u64(config.seed); + + // 1. subsample (fit on all of n only when n is already small) + let m = config.sample_cap.max(k).min(n); + + let sample = if m == n { + Cow::Borrowed(x) + } else { + let indices = sample_indices(n, m, &mut rng); + let mut sampled = vec![0_f32; m * d]; + + let chunks = sampled.chunks_mut(d); + assert_eq!(chunks.len(), indices.len()); + + for (chunk, index) in chunks.zip(indices) { + chunk.copy_from_slice(&x[index * d..(index + 1) * d]); + } + + Cow::Owned(sampled) + }; + + let sample = sample.as_ref(); + let chunk = config.chunk.get(); + let row_chunk = chunk + .checked_mul(d) + .unwrap_or_else(|| usize::MAX - (usize::MAX % d)) + .max(d); + + let sample_inv_norms: Vec = sample + .par_chunks_exact(d) + .map(|point| { + // SAFETY: every point is a `d`-sized row, and `d` is a multiple of 8 (guaranteed by + // Dimension). + let norm = unsafe { kernel::dot(point, point).sqrt() }; + + if norm > 0.0 { norm.recip() } else { 0.0 } + }) + .collect(); + + // 2. fit on the sample, best of n_init restarts (guards against bad initializations) + let mut fit = Fit::new(k, m, d); + fit.run(sample, chunk, row_chunk, &sample_inv_norms, rng, config); + mem::swap(&mut clustering.centroids, &mut fit.best_centroids); + + // 3. assign points to clusters + // SAFETY: `x.len() == n * d` (asserted above), `clustering.centroids.len() == k * d`, + // `k > 0` (checked above), `d` is a multiple of 8 (guaranteed by Dimension). + unsafe { + assign(x, &mut clustering, k, chunk, row_chunk); + } + + clustering +} + +#[cfg(test)] +mod tests { + #![expect( + clippy::float_cmp, + clippy::integer_division_remainder_used, + reason = "test module: float comparisons are intentional for exact-zero and distance \ + checks; modulo is used in test data construction" + )] + use super::*; + + /// Builds well-separated blob clusters in D-dimensional space. + /// + /// Each blob has a dominant axis so clusters are far apart in cosine + /// space. Returns `(flat_points, ground_truth_labels)`. + #[expect( + clippy::cast_possible_truncation, + reason = "k is small in tests, fits in u16" + )] + fn make_blobs( + points_per_cluster: usize, + k: usize, + seed: u64, + ) -> (Vec, Vec) { + let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); + let n = points_per_cluster * k; + let mut data = vec![0.0_f32; n * D]; + let mut truth = vec![0_u16; n]; + + for c in 0..k { + let axis = c % D; + for p in 0..points_per_cluster { + let idx = c * points_per_cluster + p; + let row = &mut data[idx * D..(idx + 1) * D]; + + row[axis] = 10.0; + for val in row.iter_mut() { + *val += rng.random_range(-0.01..0.01); + } + + truth[idx] = c as u16; + } + } + + (data, truth) + } + + const D: usize = 64; + + fn l2(v: &[f32]) -> f32 { + v.iter().map(|x| x * x).sum::().sqrt() + } + + /// Random unit-norm centroids in `D`-dimensional space. + fn unit_random(k: usize, seed: u64) -> Vec { + let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); + let mut c = vec![0.0_f32; k * D]; + for row in c.chunks_exact_mut(D) { + for v in row.iter_mut() { + *v = rng.random_range(-1.0..1.0); + } + let n = l2(row); + for v in row.iter_mut() { + *v /= n; + } + } + c + } + + /// Brute-force nearest centroid by cosine similarity. + #[expect(clippy::cast_possible_truncation, reason = "k is small in tests")] + fn brute_nearest_cosine(point: &[f32], centroids: &[f32], k: usize) -> u16 { + let pn = l2(point); + let mut best = 0_u16; + let mut best_cos = f32::NEG_INFINITY; + for c in 0..k { + let cent = ¢roids[c * D..(c + 1) * D]; + let d: f32 = point.iter().zip(cent).map(|(a, b)| a * b).sum(); + let cn = l2(cent); + let cos = if pn == 0.0 || cn == 0.0 { + 0.0 + } else { + d / (pn * cn) + }; + if cos > best_cos { + best_cos = cos; + best = c as u16; + } + } + best + } + + /// Computes clustering accuracy using majority-vote label mapping. + /// + /// K-means labels are permutation-invariant, so this assigns each + /// predicted cluster to whichever ground-truth cluster it overlaps + /// most, then counts correctly assigned points. + #[expect( + clippy::cast_precision_loss, + reason = "counts are small test values, well within f64 precision" + )] + fn accuracy(predicted: &[u16], truth: &[u16], k: usize) -> f64 { + let mut votes = vec![vec![0_usize; k]; k]; + for (&pred, &true_label) in predicted.iter().zip(truth) { + votes[pred as usize][true_label as usize] += 1; + } + + let correct: usize = votes + .iter() + .map(|row| row.iter().copied().max().unwrap_or(0)) + .sum(); + + correct as f64 / predicted.len() as f64 + } + + /// Shorthand for [`Dimension::new`] that panics on invalid input. + fn dim(d: u16) -> Dimension { + Dimension::new(d).expect("test dimension must be a positive multiple of 8") + } + + #[test] + fn chord_identical_vectors_is_zero() { + // dot=1.0, inv_norm=1.0 => similarity=1 => distance=0 + assert_eq!(squared_chord_distance(1.0, 1.0), 0.0); + } + + #[test] + fn chord_orthogonal_vectors() { + // dot=0 => similarity=0 => distance=2 + let dist = squared_chord_distance(0.0, 1.0); + assert!((dist - 2.0).abs() < 1e-6, "expected 2.0, got {dist}"); + } + + #[test] + fn chord_opposite_vectors() { + // dot=-1.0 => similarity=-1 => distance=4 + let dist = squared_chord_distance(-1.0, 1.0); + assert!((dist - 4.0).abs() < 1e-6, "expected 4.0, got {dist}"); + } + + #[test] + fn chord_zero_norm_returns_zero() { + assert_eq!(squared_chord_distance(0.5, 0.0), 0.0); + assert_eq!(squared_chord_distance(-0.5, 0.0), 0.0); + } + + #[test] + fn chord_is_non_negative() { + for dot_val in [0.0, 0.5, 1.0, -0.5, -1.0, 2.0, -2.0] { + for inv in [0.0, 0.5, 1.0, 2.0] { + let dist = squared_chord_distance(dot_val, inv); + assert!(dist >= 0.0, "negative for dot={dot_val}, inv={inv}: {dist}"); + } + } + } + + #[test] + fn cluster_empty_input() { + let config = Config::for_k_with_seed(4, 42); + let result = cluster(&[], dim(8), &config); + assert_eq!(result.labels.len(), 0); + assert_eq!(result.centroids.len(), 0); + } + + #[test] + fn cluster_k0() { + let data = vec![1.0_f32; 8]; + let config = Config::for_k_with_seed(0, 42); + let result = cluster(&data, dim(8), &config); + assert_eq!(result.labels.len(), 1); + assert_eq!(result.labels[0], 0); + } + + #[test] + fn cluster_k1_all_same_label() { + let (data, _) = make_blobs::<8>(20, 3, 123); + let config = Config::for_k_with_seed(1, 42); + let result = cluster(&data, dim(8), &config); + + assert_eq!(result.labels.len(), 60); + assert!( + result.labels.iter().all(|&l| l == 0), + "k=1: all labels must be 0" + ); + } + + #[test] + fn cluster_single_point() { + let data = vec![1.0_f32; 16]; + let config = Config::for_k_with_seed(5, 42); + // k clamped to min(k, n) = 1 + let result = cluster(&data, dim(16), &config); + assert_eq!(result.labels.len(), 1); + assert_eq!(result.labels[0], 0); + } + + #[test] + fn cluster_n_less_than_4() { + // n=3 exercises the scalar tail (no nearest4 tiling). + let (data, _) = make_blobs::<8>(1, 3, 99); + let config = Config::for_k_with_seed(3, 42); + let result = cluster(&data, dim(8), &config); + + assert_eq!(result.labels.len(), 3); + let mut seen = [false; 3]; + for &label in &*result.labels { + seen[label as usize] = true; + } + assert!( + seen.iter().all(|&s| s), + "each point should have a unique cluster" + ); + } + + #[test] + fn cluster_n_equals_k() { + let (data, _) = make_blobs::<8>(1, 5, 77); + let config = Config::for_k_with_seed(5, 42); + let result = cluster(&data, dim(8), &config); + + assert_eq!(result.labels.len(), 5); + let mut seen = [false; 5]; + for &label in &*result.labels { + seen[label as usize] = true; + } + assert!( + seen.iter().all(|&s| s), + "n=k: each point should be its own cluster" + ); + } + + #[test] + fn cluster_recovers_well_separated_blobs() { + let (data, truth) = make_blobs::<8>(50, 4, 314); + let config = Config::for_k_with_seed(4, 42); + let result = cluster(&data, dim(8), &config); + + let acc = accuracy(&result.labels, &truth, 4); + assert!( + acc > 0.95, + "expected >95% accuracy on well-separated blobs, got {:.1}%", + acc * 100.0 + ); + } + + #[test] + fn cluster_deterministic_with_same_seed() { + let (data, _) = make_blobs::<8>(30, 3, 555); + + let r1 = cluster(&data, dim(8), &Config::for_k_with_seed(3, 42)); + let r2 = cluster(&data, dim(8), &Config::for_k_with_seed(3, 42)); + + assert_eq!(r1.labels, r2.labels); + assert_eq!(r1.centroids, r2.centroids); + } + + #[test] + fn cluster_different_seeds_may_differ() { + let (data, _) = make_blobs::<8>(30, 3, 555); + + let r1 = cluster(&data, dim(8), &Config::for_k_with_seed(3, 42)); + let r2 = cluster(&data, dim(8), &Config::for_k_with_seed(3, 9999)); + + // Not guaranteed to differ, but with well-separated blobs and + // different seeds the label permutation usually differs. + assert!( + r1.labels != r2.labels, + "different seeds produced identical label vectors (possible but unlikely)" + ); + } + + #[test] + fn cluster_centroids_are_unit_normalized() { + let (data, _) = make_blobs::<8>(40, 4, 222); + let config = Config::for_k_with_seed(4, 42); + let result = cluster(&data, dim(8), &config); + + for c in 0..4_u16 { + let centroid = result.centroid(c); + // SAFETY: centroid has length 8 (= D), a multiple of 8. + let norm = unsafe { kernel::dot(centroid, centroid).sqrt() }; + assert!( + (norm - 1.0).abs() < 1e-5, + "centroid {c} has norm {norm}, expected 1.0" + ); + } + } + + #[test] + fn cluster_labels_in_range() { + let (data, _) = make_blobs::<8>(25, 5, 333); + let config = Config::for_k_with_seed(5, 42); + let result = cluster(&data, dim(8), &config); + + for (i, &label) in result.labels.iter().enumerate() { + assert!(label < 5, "label[{i}] = {label}, expected < 5"); + } + } + + #[test] + fn cluster_labels_nearest_to_assigned_centroid() { + let (data, _) = make_blobs::<8>(30, 3, 444); + let config = Config::for_k_with_seed(3, 42); + let result = cluster(&data, dim(8), &config); + + let k = 3_usize; + let d = 8_usize; + for (i, point) in data.chunks_exact(d).enumerate() { + let assigned = result.labels[i]; + // SAFETY: point and centroid both have length 8 (= D), a multiple of 8. + let assigned_dot = unsafe { kernel::dot(point, result.centroid(assigned)) }; + + #[expect(clippy::cast_possible_truncation, reason = "k=3 fits in u16")] + for c in 0..k as u16 { + // SAFETY: point and centroid both have length 8, a multiple of 8. + let other_dot = unsafe { kernel::dot(point, result.centroid(c)) }; + assert!( + other_dot <= assigned_dot + 1e-5, + "point {i}: assigned to {assigned} (dot={assigned_dot}) but centroid {c} has \ + higher dot={other_dot}" + ); + } + } + } + + #[test] + fn cluster_d32_recovers_blobs() { + let (data, truth) = make_blobs::<32>(40, 3, 888); + let config = Config::for_k_with_seed(3, 42); + let result = cluster(&data, dim(32), &config); + + let acc = accuracy(&result.labels, &truth, 3); + assert!( + acc > 0.95, + "D=32: expected >95% accuracy, got {:.1}%", + acc * 100.0 + ); + } + + #[test] + fn cluster_recovers_with_subsampling() { + // n=12000 with sample_cap=1024 exercises the Cow::Owned path. + let (data, truth) = make_blobs::<8>(2000, 6, 21); + let mut config = Config::for_k_with_seed(6, 5); + config.sample_cap = 1024; + let result = cluster(&data, dim(8), &config); + + let acc = accuracy(&result.labels, &truth, 6); + assert!( + acc > 0.95, + "subsampled: expected >95% accuracy, got {:.1}%", + acc * 100.0 + ); + } + + #[test] + fn cluster_more_clusters_than_natural_groups() { + // 3 natural groups but k=8: empty clusters keep their seed centroid, + // nothing should be NaN or infinite. + let (data, _) = make_blobs::<8>(400, 3, 31); + let result = cluster(&data, dim(8), &Config::for_k_with_seed(8, 1)); + + assert!( + result.centroids.iter().all(|v| v.is_finite()), + "NaN or infinite centroid" + ); + assert!(result.labels.iter().all(|&l| l < 8)); + } + + #[test] + fn cluster_all_identical_points() { + // Every point identical: D² distances are all zero during seeding, + // which triggers the uniform fallback path. + let n = 100; + let mut data = vec![0.0_f32; n * 8]; + for row in data.chunks_exact_mut(8) { + row[0] = 1.0; + } + let result = cluster(&data, dim(8), &Config::for_k_with_seed(4, 1)); + + assert!(result.centroids.iter().all(|v| v.is_finite())); + assert!(result.labels.iter().all(|&l| l < 4)); + } + + #[test] + fn nearest_centroid_matches_brute_force_cosine() { + let k = 7; + let centroids = unit_random(k, 99); + let mut rng = Xoshiro256PlusPlus::seed_from_u64(100); + + for _ in 0..1000 { + let p: Vec = core::iter::repeat_with(|| rng.random_range(-3.0..3.0)) + .take(D) + .collect(); + let pn = l2(&p); + let inv = if pn > 0.0 { pn.recip() } else { 0.0 }; + + // SAFETY: point has length D=64, centroids has length k*D, + // k > 0, D is a multiple of 8. + let (got, _) = unsafe { nearest_centroid(&p, inv, ¢roids, k, D) }; + assert_eq!( + got, + brute_nearest_cosine(&p, ¢roids, k), + "mismatch for point norm={pn}" + ); + } + } + + #[test] + fn nearest_centroid_argmax_independent_of_inv_norm() { + let k = 5; + let centroids = unit_random(k, 7); + let mut rng = Xoshiro256PlusPlus::seed_from_u64(8); + + for _ in 0..500 { + let p: Vec = core::iter::repeat_with(|| rng.random_range(-2.0..2.0)) + .take(D) + .collect(); + + // SAFETY: point has length D=64, centroids has length k*D, + // k > 0, D is a multiple of 8. + let (a, _) = unsafe { nearest_centroid(&p, 1.0, ¢roids, k, D) }; + // SAFETY: same preconditions. + let (b, _) = unsafe { nearest_centroid(&p, 0.123, ¢roids, k, D) }; + assert_eq!(a, b, "inv_norm must not change the selected centroid"); + } + } + + #[test] + fn cluster_mixed_zero_norm_rows() { + // Some all-zero rows exercise the inv_norm == 0 path in accumulation + // and the squared_chord_distance == 0 return. + let n = 120; + let mut data = vec![0.0_f32; n * 8]; + let mut rng = Xoshiro256PlusPlus::seed_from_u64(7); + for (i, row) in data.chunks_exact_mut(8).enumerate() { + if i % 10 == 0 { + continue; // leave all-zero + } + for v in row.iter_mut() { + *v = rng.random_range(-1.0..1.0); + } + } + let result = cluster(&data, dim(8), &Config::for_k_with_seed(5, 1)); + + assert!(result.centroids.iter().all(|v| v.is_finite())); + assert!(result.labels.iter().all(|&l| l < 5)); + } +} diff --git a/libs/@local/graph/store/src/embedding/dimension.rs b/libs/@local/graph/store/src/embedding/dimension.rs new file mode 100644 index 00000000000..a3a1a31fe94 --- /dev/null +++ b/libs/@local/graph/store/src/embedding/dimension.rs @@ -0,0 +1,79 @@ +use core::num::NonZero; + +/// An embedding vector dimension, guaranteed to be a positive multiple of 8. +/// +/// The multiple-of-8 invariant ensures that the dimension evenly divides into +/// SIMD lanes (8×f32 = `f32x8`), so vectorized kernels can operate without +/// remainder handling. +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Dimension(NonZero); + +impl Dimension { + /// Creates a new dimension if `value` is non-zero and a multiple of 8. + /// + /// Returns [`None`] otherwise. + #[must_use] + pub const fn new(value: u16) -> Option { + // not using `?` here because it isn't `const` + let Some(value) = NonZero::new(value) else { + return None; + }; + + if !value.get().is_multiple_of(8) { + return None; + } + + Some(Self(value)) + } + + /// The raw dimension value. + #[must_use] + pub const fn get(self) -> u16 { + self.0.get() + } +} + +pub const D128: Dimension = Dimension(NonZero::new(128).unwrap()); +pub const D256: Dimension = Dimension(NonZero::new(256).unwrap()); +pub const D512: Dimension = Dimension(NonZero::new(512).unwrap()); +pub const D1536: Dimension = Dimension(NonZero::new(1536).unwrap()); +pub const D3072: Dimension = Dimension(NonZero::new(3072).unwrap()); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn valid_multiples_of_8() { + for v in [8, 16, 24, 128, 256, 3072] { + assert!( + Dimension::new(v).is_some(), + "{v} should be a valid dimension" + ); + } + } + + #[test] + fn zero_rejected() { + assert!(Dimension::new(0).is_none()); + } + + #[test] + fn non_multiples_of_8_rejected() { + for v in [1, 2, 3, 4, 5, 6, 7, 9, 10, 15, 17, 100, 3071] { + assert!( + Dimension::new(v).is_none(), + "{v} should not be a valid dimension" + ); + } + } + + #[test] + fn constants_have_correct_values() { + assert_eq!(D128.0.get(), 128); + assert_eq!(D256.0.get(), 256); + assert_eq!(D512.0.get(), 512); + assert_eq!(D1536.0.get(), 1536); + assert_eq!(D3072.0.get(), 3072); + } +} diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/store/src/embedding/kernel.rs new file mode 100644 index 00000000000..0090f63c4c5 --- /dev/null +++ b/libs/@local/graph/store/src/embedding/kernel.rs @@ -0,0 +1,769 @@ +use core::simd::{Simd, f32x8, num::SimdFloat as _}; + +/// Fused multiply-add when the target has native FMA, separate mul+add otherwise. +/// +/// On aarch64, FMA is part of the base NEON instruction set (`fmla`). +/// On x86_64, FMA requires the `fma` target feature (`vfmadd`); without it, +/// `StdFloat::mul_add` falls back to a per-lane `fmaf` libc call which +/// destroys throughput. The non-FMA path uses a plain multiply and add +/// (`vmulps` + `vaddps`) instead. +#[inline(always)] +#[cfg(not(any(target_arch = "aarch64", target_feature = "fma")))] +fn simd_mul_add(lhs: f32x8, rhs: f32x8, acc: f32x8) -> f32x8 { + lhs * rhs + acc +} + +/// See non-FMA variant above for rationale. +#[inline(always)] +#[cfg(any(target_arch = "aarch64", target_feature = "fma"))] +fn simd_mul_add(lhs: f32x8, rhs: f32x8, acc: f32x8) -> f32x8 { + use std::simd::StdFloat as _; + + lhs.mul_add(rhs, acc) +} + +/// Computes the dot product of two equal-length `f32` slices using SIMD. +/// +/// Four independent accumulators are interleaved to saturate FMA throughput: +/// each accumulator feeds a separate dependency chain, hiding the 4-cycle +/// latency of `fmla`/`vfmadd` on typical micro-architectures. +/// +/// # Safety +/// +/// * `lhs.len() == rhs.len()` +/// * Both lengths are multiples of 8. +#[inline] +#[must_use] +pub(crate) unsafe fn dot(lhs: &[f32], rhs: &[f32]) -> f32 { + debug_assert!(lhs.len().is_multiple_of(8) && lhs.len() == rhs.len()); + + // SAFETY: the caller guarantees equal lengths and a multiple of 8. + // These hints let the compiler elide bounds checks in `as_chunks` and + // the subsequent indexing without raw pointer arithmetic. + unsafe { + core::hint::assert_unchecked(lhs.len() == rhs.len()); + core::hint::assert_unchecked(lhs.len().is_multiple_of(8)); + } + + let (lhs, _) = lhs.as_chunks::<8>(); + let (rhs, _) = rhs.as_chunks::<8>(); + + // SAFETY: both original slices have the same length and that length is a + // multiple of 8, so `as_chunks::<8>` produces equal-length chunk slices + // with empty remainders. + unsafe { + core::hint::assert_unchecked(lhs.len() == rhs.len()); + } + + let mut s0 = f32x8::splat(0.0); + let mut s1 = f32x8::splat(0.0); + let mut s2 = f32x8::splat(0.0); + let mut s3 = f32x8::splat(0.0); + + // Unrolled loop: process 4 chunks (32 floats) per iteration. + let mut offset = 0; + while offset + 4 <= lhs.len() { + let Some([l0, l1, l2, l3]) = lhs[offset..offset + 4].as_array() else { + unreachable!() + }; + let Some([r0, r1, r2, r3]) = rhs[offset..offset + 4].as_array() else { + unreachable!() + }; + + s0 = simd_mul_add(Simd::from_slice(l0), Simd::from_slice(r0), s0); + s1 = simd_mul_add(Simd::from_slice(l1), Simd::from_slice(r1), s1); + s2 = simd_mul_add(Simd::from_slice(l2), Simd::from_slice(r2), s2); + s3 = simd_mul_add(Simd::from_slice(l3), Simd::from_slice(r3), s3); + + offset += 4; + } + + // Tail: process remaining 0..3 chunks one at a time. + #[expect(clippy::min_ident_chars)] + while offset < lhs.len() { + let l = &lhs[offset]; + let r = &rhs[offset]; + + s0 = simd_mul_add(Simd::from_slice(l), Simd::from_slice(r), s0); + offset += 1; + } + + (s0 + s1 + s2 + s3).reduce_sum() +} + +/// Adds `src` element-wise into `dst`. +/// +/// # Safety +/// +/// * `dst.len() == src.len()` +/// * Both lengths are multiples of 8. +#[inline] +pub(crate) unsafe fn add_into(dst: &mut [f32], src: &[f32]) { + debug_assert!(dst.len().is_multiple_of(8) && dst.len() == src.len()); + + // SAFETY: the caller guarantees equal lengths and a multiple of 8. + unsafe { + core::hint::assert_unchecked(dst.len() == src.len()); + core::hint::assert_unchecked(dst.len().is_multiple_of(8)); + } + + let (dst, _) = dst.as_chunks_mut::<8>(); + let (src, _) = src.as_chunks::<8>(); + + // SAFETY: same reasoning as the pre-chunk hints: equal input lengths + // that are multiples of 8 produce equal chunk counts. + unsafe { core::hint::assert_unchecked(dst.len() == src.len()) } + + for index in 0..dst.len() { + dst[index] = (f32x8::from_slice(&dst[index]) + f32x8::from_slice(&src[index])).to_array(); + } +} + +/// Writes `src * factor` element-wise into `dst`. +/// +/// # Safety +/// +/// * `dst.len() == src.len()` +/// * Both lengths are multiples of 8. +#[inline] +pub(crate) unsafe fn scale_into(dst: &mut [f32], src: &[f32], factor: f32) { + debug_assert!(dst.len().is_multiple_of(8) && dst.len() == src.len()); + + // SAFETY: the caller guarantees equal lengths and a multiple of 8. + unsafe { + core::hint::assert_unchecked(dst.len() == src.len()); + core::hint::assert_unchecked(dst.len().is_multiple_of(8)); + } + + let factor = f32x8::splat(factor); + let (dst, _) = dst.as_chunks_mut::<8>(); + let (src, _) = src.as_chunks::<8>(); + + // SAFETY: same reasoning as the pre-chunk hints: equal input lengths + // that are multiples of 8 produce equal chunk counts. + unsafe { core::hint::assert_unchecked(dst.len() == src.len()) } + + for index in 0..dst.len() { + dst[index] = (f32x8::from_slice(&src[index]) * factor).to_array(); + } +} + +/// Scales `value` in-place by `factor`. +/// +/// # Safety +/// +/// * `value.len()` is a multiple of 8. +#[inline] +pub(crate) unsafe fn scale(value: &mut [f32], factor: f32) { + debug_assert!(value.len().is_multiple_of(8)); + + // SAFETY: the caller guarantees a multiple of 8. + unsafe { + core::hint::assert_unchecked(value.len().is_multiple_of(8)); + } + + let factor = f32x8::splat(factor); + let (dst, _) = value.as_chunks_mut::<8>(); + + for dst in dst { + *dst = (f32x8::from_slice(dst) * factor).to_array(); + } +} + +/// Accumulates `src * factor` element-wise into `dst` (`dst += src * factor`). +/// +/// Fuses a scale and add into a single pass, using FMA where available. +/// Avoids the need for a scratch buffer when accumulating normalized vectors. +/// +/// # Safety +/// +/// * `dst.len() == src.len()` +/// * Both lengths are multiples of 8. +#[inline] +pub(crate) unsafe fn add_scaled_into(dst: &mut [f32], src: &[f32], factor: f32) { + debug_assert!(dst.len().is_multiple_of(8) && dst.len() == src.len()); + + // SAFETY: the caller guarantees equal lengths and a multiple of 8. + unsafe { + core::hint::assert_unchecked(dst.len() == src.len()); + core::hint::assert_unchecked(dst.len().is_multiple_of(8)); + } + + let factor = f32x8::splat(factor); + let (dst, _) = dst.as_chunks_mut::<8>(); + let (src, _) = src.as_chunks::<8>(); + + // SAFETY: same reasoning as the pre-chunk hints: equal input lengths + // that are multiples of 8 produce equal chunk counts. + unsafe { core::hint::assert_unchecked(dst.len() == src.len()) } + + for index in 0..dst.len() { + let acc = f32x8::from_slice(&dst[index]); + let val = f32x8::from_slice(&src[index]); + dst[index] = simd_mul_add(val, factor, acc).to_array(); + } +} + +/// Normalizes `value` to unit length in-place. +/// +/// If the vector has zero norm, it is left unchanged. +/// +/// # Safety +/// +/// * `value.len()` is a multiple of 8. +#[inline] +pub(crate) unsafe fn normalize(value: &mut [f32]) { + // SAFETY: `dot` requires equal lengths (trivially true, same slice) + // and a multiple of 8 (guaranteed by the caller). + let norm = unsafe { dot(value, value).sqrt() }; + + if norm > 0.0 { + let factor = 1.0 / norm; + // SAFETY: same slice, same length guarantee. + unsafe { + scale(value, factor); + } + } +} + +/// 4 points x 2 centroids. Eight independent accumulators give ILP 8 (enough to +/// saturate FMA throughput); each point chunk feeds 2 FMAs and each centroid +/// chunk feeds 4. Returns `dot[point][centroid]`. +/// +/// Register budget: 8 `f32x8` accumulators. On AVX2 (16 ymm) this leaves room +/// for the 6 operand loads. On NEON each `f32x8` is two 128-bit regs, so the 8 +/// accumulators take 16 of 32 registers; a 4x4 tile (16 accumulators) also fits +/// there if you want more centroid reuse. Either way, check the asm shows the +/// accumulators staying in registers with no stack spills, and that +/// `simd_mul_add` lowered to `vfmadd`/`fmla` and not a `fmaf` call. If the array +/// form ever spills, the manual unroll below is what keeps them in registers. +/// +/// # Safety +/// * all six slices have length `d` +/// * `d` is a multiple of 8 +#[expect( + clippy::inline_always, + reason = "micro-kernel must inline into nearest4 to keep accumulators in registers" +)] +#[inline(always)] +pub(crate) unsafe fn micro_4x2( + p0: &[f32], + p1: &[f32], + p2: &[f32], + p3: &[f32], + c0: &[f32], + c1: &[f32], +) -> [[f32; 2]; 4] { + debug_assert!(p0.len().is_multiple_of(8)); + debug_assert!( + [p1.len(), p2.len(), p3.len(), c0.len(), c1.len()] + .iter() + .all(|&l| l == p0.len()) + ); + + let (p0, _) = p0.as_chunks::<8>(); + let (p1, _) = p1.as_chunks::<8>(); + let (p2, _) = p2.as_chunks::<8>(); + let (p3, _) = p3.as_chunks::<8>(); + let (c0, _) = c0.as_chunks::<8>(); + let (c1, _) = c1.as_chunks::<8>(); + + // SAFETY: the caller guarantees all six slices have equal length `d`, + // and `d` is a multiple of 8. The hints let the compiler prove that + // `as_chunks` produces equal-length chunk slices. + unsafe { + core::hint::assert_unchecked(p0.len() == p1.len()); + core::hint::assert_unchecked(p0.len() == p2.len()); + core::hint::assert_unchecked(p0.len() == p3.len()); + core::hint::assert_unchecked(p0.len() == c0.len()); + core::hint::assert_unchecked(p0.len() == c1.len()); + } + + let mut a00 = f32x8::splat(0.0); + let mut a01 = f32x8::splat(0.0); + let mut a10 = f32x8::splat(0.0); + let mut a11 = f32x8::splat(0.0); + let mut a20 = f32x8::splat(0.0); + let mut a21 = f32x8::splat(0.0); + let mut a30 = f32x8::splat(0.0); + let mut a31 = f32x8::splat(0.0); + + for t in 0..c0.len() { + let v0 = Simd::from_array(c0[t]); + let v1 = Simd::from_array(c1[t]); + let x0 = Simd::from_array(p0[t]); + let x1 = Simd::from_array(p1[t]); + let x2 = Simd::from_array(p2[t]); + let x3 = Simd::from_array(p3[t]); + + // super::simd_mul_add picks the FMA arm per target. + a00 = simd_mul_add(x0, v0, a00); + a01 = simd_mul_add(x0, v1, a01); + a10 = simd_mul_add(x1, v0, a10); + a11 = simd_mul_add(x1, v1, a11); + a20 = simd_mul_add(x2, v0, a20); + a21 = simd_mul_add(x2, v1, a21); + a30 = simd_mul_add(x3, v0, a30); + a31 = simd_mul_add(x3, v1, a31); + } + + [ + [a00.reduce_sum(), a01.reduce_sum()], + [a10.reduce_sum(), a11.reduce_sum()], + [a20.reduce_sum(), a21.reduce_sum()], + [a30.reduce_sum(), a31.reduce_sum()], + ] +} + +/// Finds the nearest centroid for 4 points simultaneously using the +/// [`micro_4x2`] tiled kernel. +/// +/// Returns `(centroid_index, raw_dot_product)` for each of the 4 points. +/// The raw dot product is **not** a distance; the caller must convert via +/// [`squared_chord_distance`](super::clustering::squared_chord_distance) +/// if needed. +/// +/// # Safety +/// +/// * All four point slices have length `d`. +/// * `centroids.len() >= k * d`. +/// * `d` is a multiple of 8. +/// * `k > 0`. +#[inline] +#[must_use] +pub(crate) unsafe fn nearest4( + p0: &[f32], + p1: &[f32], + p2: &[f32], + p3: &[f32], + centroids: &[f32], + k: usize, + d: usize, +) -> [(u16, f32); 4] { + let mut best_dot = [f32::NEG_INFINITY; 4]; + let mut best_idx = [0_u16; 4]; + + // SAFETY: the caller guarantees these preconditions. + unsafe { + core::hint::assert_unchecked(p0.len() == d); + core::hint::assert_unchecked(p0.len() == p1.len()); + core::hint::assert_unchecked(p0.len() == p2.len()); + core::hint::assert_unchecked(p0.len() == p3.len()); + core::hint::assert_unchecked(centroids.len() >= k * d); + core::hint::assert_unchecked(d.is_multiple_of(8)); + core::hint::assert_unchecked(k > 0); + } + + let mut j = 0; + while j + 2 <= k { + // SAFETY: `j + 2 <= k` and `centroids.len() >= k * d`, so both + // slices `[j*d .. (j+2)*d]` are in-bounds. + let c0 = unsafe { centroids.get_unchecked(j * d..j * d + d) }; + // SAFETY: see above. + let c1 = unsafe { centroids.get_unchecked((j + 1) * d..(j + 1) * d + d) }; + + // SAFETY: all six slices have length `d`, a multiple of 8. + let dots = unsafe { micro_4x2(p0, p1, p2, p3, c0, c1) }; + + #[expect( + clippy::cast_possible_truncation, + reason = "k originates from Config::k (u16), so j < k fits in u16" + )] + for m in 0..4 { + if dots[m][0] > best_dot[m] { + best_dot[m] = dots[m][0]; + best_idx[m] = j as u16; + } + if dots[m][1] > best_dot[m] { + best_dot[m] = dots[m][1]; + best_idx[m] = (j + 1) as u16; + } + } + j += 2; + } + + // Handle odd k: one remaining centroid. + if j < k { + let c = ¢roids[j * d..j * d + d]; + let ps = [p0, p1, p2, p3]; + for m in 0..4 { + // SAFETY: point and centroid both have length `d`, a multiple of 8. + let d = unsafe { dot(ps[m], c) }; + #[expect( + clippy::cast_possible_truncation, + reason = "k originates from Config::k (u16)" + )] + if d > best_dot[m] { + best_dot[m] = d; + best_idx[m] = j as u16; + } + } + } + + [ + (best_idx[0], best_dot[0]), + (best_idx[1], best_dot[1]), + (best_idx[2], best_dot[2]), + (best_idx[3], best_dot[3]), + ] +} + +#[cfg(test)] +mod tests { + #![expect(clippy::float_cmp, clippy::integer_division_remainder_used)] + + use super::*; + + /// Scalar dot product for reference. + fn ref_dot(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| x * y).sum() + } + + /// Scalar normalize for reference. + fn ref_normalize(v: &mut [f32]) { + let norm = v.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for x in v { + *x /= norm; + } + } + } + + /// Deterministic test vector: entry `i` gets `(i+1) * scale`. + #[expect(clippy::cast_precision_loss)] + fn ramp(len: usize, factor: f32) -> Vec { + (0..len).map(|i| (i + 1) as f32 * factor).collect() + } + + /// Asserts two f32 values are within relative tolerance, with an absolute + /// floor for values near zero. + fn assert_close(a: f32, b: f32, tol: f32) { + let diff = (a - b).abs(); + let denom = a.abs().max(b.abs()).max(1e-12); + assert!( + diff / denom < tol, + "values differ: {a} vs {b} (diff={diff}, rel={})", + diff / denom + ); + } + + #[test] + fn dot_matches_scalar_d8() { + let a = ramp(8, 1.0); + let b = ramp(8, 0.5); + let expected = ref_dot(&a, &b); + // SAFETY: both slices have length 8, a multiple of 8. + let got = unsafe { dot(&a, &b) }; + assert_close(got, expected, 1e-6); + } + + #[test] + fn dot_matches_scalar_d24() { + // 24 = 3 chunks of 8: the 4-unrolled body runs 0 iterations, + // all 3 chunks go through the tail path. + let a = ramp(24, 0.1); + let b = ramp(24, -0.2); + let expected = ref_dot(&a, &b); + // SAFETY: both slices have length 24, a multiple of 8. + let got = unsafe { dot(&a, &b) }; + assert_close(got, expected, 1e-5); + } + + #[test] + fn dot_matches_scalar_d3072() { + let a = ramp(3072, 0.001); + let b = ramp(3072, -0.002); + let expected = ref_dot(&a, &b); + // SAFETY: both slices have length 3072, a multiple of 8. + let got = unsafe { dot(&a, &b) }; + assert_close(got, expected, 1e-4); + } + + #[test] + fn dot_is_commutative() { + let a = ramp(32, 0.3); + let b = ramp(32, -0.7); + // SAFETY: both slices have length 32, a multiple of 8. + let ab = unsafe { dot(&a, &b) }; + // SAFETY: same slices, reversed. + let ba = unsafe { dot(&b, &a) }; + assert_eq!(ab, ba); + } + + #[test] + fn dot_self_is_squared_norm() { + let a = ramp(16, 0.5); + let expected: f32 = a.iter().map(|x| x * x).sum(); + // SAFETY: both arguments are the same 16-element slice. + let got = unsafe { dot(&a, &a) }; + assert_close(got, expected, 1e-6); + } + + #[test] + fn dot_orthogonal_is_zero() { + let mut a = vec![0.0_f32; 8]; + let mut b = vec![0.0_f32; 8]; + a[0] = 1.0; + b[1] = 1.0; + // SAFETY: both slices have length 8. + let got = unsafe { dot(&a, &b) }; + assert_eq!(got, 0.0); + } + + #[test] + fn add_into_matches_scalar() { + let src = ramp(16, 1.0); + let mut dst = ramp(16, 0.5); + let expected: Vec = dst.iter().zip(&src).map(|(d, s)| d + s).collect(); + // SAFETY: both slices have length 16, a multiple of 8. + unsafe { add_into(&mut dst, &src) } + assert_eq!(dst, expected); + } + + #[test] + fn add_into_zero_is_identity() { + let zeros = vec![0.0_f32; 24]; + let mut dst = ramp(24, 1.0); + let original = dst.clone(); + // SAFETY: both slices have length 24, a multiple of 8. + unsafe { add_into(&mut dst, &zeros) } + assert_eq!(dst, original); + } + + #[test] + fn scale_into_matches_scalar() { + let src = ramp(16, 1.0); + let mut dst = vec![0.0_f32; 16]; + let factor = 2.5; + let expected: Vec = src.iter().map(|x| x * factor).collect(); + // SAFETY: both slices have length 16, a multiple of 8. + unsafe { scale_into(&mut dst, &src, factor) } + assert_eq!(dst, expected); + } + + #[test] + fn scale_into_zero_gives_zeros() { + let src = ramp(8, 1.0); + let mut dst = ramp(8, 999.0); + // SAFETY: both slices have length 8, a multiple of 8. + unsafe { scale_into(&mut dst, &src, 0.0) } + assert!(dst.iter().all(|&x| x == 0.0)); + } + + #[test] + fn scale_into_one_is_copy() { + let src = ramp(16, 0.3); + let mut dst = vec![0.0_f32; 16]; + // SAFETY: both slices have length 16, a multiple of 8. + unsafe { scale_into(&mut dst, &src, 1.0) } + assert_eq!(dst, src); + } + + #[test] + fn scale_matches_scalar() { + let mut v = ramp(16, 1.0); + let factor = -0.5; + let expected: Vec = v.iter().map(|x| x * factor).collect(); + // SAFETY: slice has length 16, a multiple of 8. + unsafe { scale(&mut v, factor) } + assert_eq!(v, expected); + } + + #[test] + fn add_scaled_into_matches_separate_ops() { + let src = ramp(16, 1.0); + let factor = 0.3; + let mut dst = ramp(16, 0.5); + let expected: Vec = dst.iter().zip(&src).map(|(d, s)| d + s * factor).collect(); + + // SAFETY: both slices have length 16, a multiple of 8. + unsafe { add_scaled_into(&mut dst, &src, factor) } + + for (&got, &exp) in dst.iter().zip(&expected) { + assert_close(got, exp, 1e-6); + } + } + + #[test] + fn add_scaled_into_factor_zero_is_identity() { + let src = ramp(8, 100.0); + let mut dst = ramp(8, 1.0); + let original = dst.clone(); + // SAFETY: both slices have length 8, a multiple of 8. + unsafe { add_scaled_into(&mut dst, &src, 0.0) } + assert_eq!(dst, original); + } + + #[test] + fn normalize_produces_unit_norm() { + let mut v = ramp(32, 0.7); + // SAFETY: length 32, a multiple of 8. + unsafe { normalize(&mut v) } + // SAFETY: same slice, length unchanged. + let norm = unsafe { dot(&v, &v).sqrt() }; + assert_close(norm, 1.0, 1e-6); + } + + #[test] + fn normalize_preserves_direction() { + let mut v = ramp(16, 2.0); + let mut ref_v = v.clone(); + ref_normalize(&mut ref_v); + // SAFETY: length 16, a multiple of 8. + unsafe { normalize(&mut v) } + for (&a, &b) in v.iter().zip(&ref_v) { + assert_close(a, b, 1e-6); + } + } + + #[test] + fn normalize_zero_vector_unchanged() { + let mut v = vec![0.0_f32; 8]; + // SAFETY: length 8, a multiple of 8. + unsafe { normalize(&mut v) } + assert!(v.iter().all(|&x| x == 0.0)); + } + + #[test] + fn normalize_already_unit_is_stable() { + let mut v = vec![0.0_f32; 8]; + v[0] = 1.0; + // SAFETY: length 8, a multiple of 8. + unsafe { normalize(&mut v) } + assert_close(v[0], 1.0, 1e-7); + assert!(v[1..].iter().all(|&x| x == 0.0)); + } + + #[test] + fn micro_4x2_matches_individual_dots() { + let d = 16; + let p0 = ramp(d, 0.1); + let p1 = ramp(d, -0.2); + let p2 = ramp(d, 0.3); + let p3 = ramp(d, -0.4); + let c0 = ramp(d, 0.5); + let c1 = ramp(d, -0.6); + + // SAFETY: all 6 slices have length 16, a multiple of 8. + let got = unsafe { micro_4x2(&p0, &p1, &p2, &p3, &c0, &c1) }; + + let expected = [ + [ref_dot(&p0, &c0), ref_dot(&p0, &c1)], + [ref_dot(&p1, &c0), ref_dot(&p1, &c1)], + [ref_dot(&p2, &c0), ref_dot(&p2, &c1)], + [ref_dot(&p3, &c0), ref_dot(&p3, &c1)], + ]; + + for (g, e) in got.iter().zip(&expected) { + assert_close(g[0], e[0], 1e-5); + assert_close(g[1], e[1], 1e-5); + } + } + + #[test] + fn micro_4x2_d3072() { + let d = 3072; + let p0 = ramp(d, 0.001); + let p1 = ramp(d, -0.001); + let p2 = ramp(d, 0.002); + let p3 = ramp(d, -0.002); + let c0 = ramp(d, 0.001); + let c1 = ramp(d, -0.001); + + // SAFETY: all 6 slices have length 3072, a multiple of 8. + let got = unsafe { micro_4x2(&p0, &p1, &p2, &p3, &c0, &c1) }; + + let expected = [ + [ref_dot(&p0, &c0), ref_dot(&p0, &c1)], + [ref_dot(&p1, &c0), ref_dot(&p1, &c1)], + [ref_dot(&p2, &c0), ref_dot(&p2, &c1)], + [ref_dot(&p3, &c0), ref_dot(&p3, &c1)], + ]; + + for (g, e) in got.iter().zip(&expected) { + assert_close(g[0], e[0], 1e-3); + assert_close(g[1], e[1], 1e-3); + } + } + + #[test] + fn nearest4_matches_brute_force_even_k() { + let d = 8; + let k = 4; + + // 4 centroids: axis-aligned unit vectors. + let mut centroids = vec![0.0_f32; k * d]; + for i in 0..k { + centroids[i * d + i] = 1.0; + } + + // 4 points, each close to a different centroid. + let mut points: [Vec; 4] = core::array::from_fn(|_| vec![0.0_f32; d]); + for i in 0..4 { + points[i][i] = 10.0; + points[i][(i + 1) % d] = 0.1; + } + + // SAFETY: d=8 (multiple of 8), k=4 > 0, centroids has length k*d, + // all point slices have length d. + let got = unsafe { + nearest4( + &points[0], &points[1], &points[2], &points[3], ¢roids, k, d, + ) + }; + + assert_eq!(got[0].0, 0); + assert_eq!(got[1].0, 1); + assert_eq!(got[2].0, 2); + assert_eq!(got[3].0, 3); + } + + #[test] + fn nearest4_matches_brute_force_odd_k() { + let d = 8; + let k = 3; // odd: exercises the remainder path + + let mut centroids = vec![0.0_f32; k * d]; + for i in 0..k { + centroids[i * d + i] = 1.0; + } + + let mut points: [Vec; 4] = core::array::from_fn(|_| vec![0.0_f32; d]); + points[0][0] = 5.0; + points[1][1] = 5.0; + points[2][2] = 5.0; + points[3][0] = 3.0; // closest to centroid 0 + + // SAFETY: d=8 (multiple of 8), k=3 > 0, centroids has length k*d, + // all point slices have length d. + let got = unsafe { + nearest4( + &points[0], &points[1], &points[2], &points[3], ¢roids, k, d, + ) + }; + + assert_eq!(got[0].0, 0); + assert_eq!(got[1].0, 1); + assert_eq!(got[2].0, 2); + assert_eq!(got[3].0, 0); + } + + #[test] + fn nearest4_k1_all_same() { + let d = 8; + let centroids = ramp(d, 1.0); + let p0 = ramp(d, 0.1); + let p1 = ramp(d, -0.2); + let p2 = ramp(d, 0.3); + let p3 = ramp(d, -0.4); + + // SAFETY: d=8 (multiple of 8), k=1 > 0, centroids has length d, + // all point slices have length d. + let got = unsafe { nearest4(&p0, &p1, &p2, &p3, ¢roids, 1, d) }; + + assert_eq!(got[0].0, 0); + assert_eq!(got[1].0, 0); + assert_eq!(got[2].0, 0); + assert_eq!(got[3].0, 0); + } +} diff --git a/libs/@local/graph/store/src/embedding/mod.rs b/libs/@local/graph/store/src/embedding/mod.rs new file mode 100644 index 00000000000..cefb998fb67 --- /dev/null +++ b/libs/@local/graph/store/src/embedding/mod.rs @@ -0,0 +1,15 @@ +#![expect( + unsafe_code, + dead_code, + clippy::indexing_slicing, + clippy::float_arithmetic, + clippy::min_ident_chars, + clippy::many_single_char_names, + reason = "embedding module is under active development; dead_code is expected until the \ + public API is wired up. Single-char idents (k, n, m, d, x) are standard \ + mathematical notation for clustering." +)] + +pub mod clustering; +pub mod dimension; +pub(crate) mod kernel; diff --git a/libs/@local/graph/store/src/entity/mod.rs b/libs/@local/graph/store/src/entity/mod.rs index f1174879703..df563a365cc 100644 --- a/libs/@local/graph/store/src/entity/mod.rs +++ b/libs/@local/graph/store/src/entity/mod.rs @@ -4,14 +4,14 @@ pub use self::{ EntityQuerySortingToken, EntityQueryToken, }, store::{ - ClosedMultiEntityTypeMap, CreateEntityParams, DeleteEntitiesParams, DeletionScope, - DeletionSummary, DiffEntityParams, DiffEntityResult, EntityPermissions, EntityStore, - EntityValidationType, HasPermissionForEntitiesParams, LinkDeletionBehavior, - PatchEntityParams, QueryConversion, QueryEntitiesParams, QueryEntitiesResponse, - QueryEntitySubgraphParams, QueryEntitySubgraphResponse, SearchEntitiesFilter, - SearchEntitiesParams, SearchEntitiesResponse, SummarizeEntitiesParams, - SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, ValidateEntityComponents, - ValidateEntityError, ValidateEntityParams, + ClosedMultiEntityTypeMap, ClusterEntitiesParams, ClusterEntitiesResponse, + CreateEntityParams, DeleteEntitiesParams, DeletionScope, DeletionSummary, DiffEntityParams, + DiffEntityResult, EntityCluster, EntityPermissions, EntityStore, EntityValidationType, + HasPermissionForEntitiesParams, LinkDeletionBehavior, PatchEntityParams, QueryConversion, + QueryEntitiesParams, QueryEntitiesResponse, QueryEntitySubgraphParams, + QueryEntitySubgraphResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, + UpdateEntityEmbeddingsParams, ValidateEntityComponents, ValidateEntityError, + ValidateEntityParams, }, validation_report::{ EmptyEntityTypes, EntityRetrieval, EntityTypeRetrieval, EntityTypesError, diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 94ff152a36b..1af2a5af1c4 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -36,7 +36,9 @@ use utoipa::{ use crate::{ entity::{EntityQueryCursor, EntityQuerySorting, EntityValidationReport}, entity_type::{EntityTypeResolveDefinitions, IncludeEntityTypeOption}, - error::{CheckPermissionError, DeletionError, InsertionError, QueryError, UpdateError}, + error::{ + CheckPermissionError, ClusterError, DeletionError, InsertionError, QueryError, UpdateError, + }, filter::{Filter, SemanticDistance}, subgraph::{ Subgraph, @@ -525,6 +527,55 @@ impl PatchEntityParams { } } +#[derive(Debug, Deserialize)] +#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ClusterEntitiesParams { + pub entity_ids: Vec, + /// Desired number of clusters. Clamped to the number of entities with + /// embeddings when that is smaller. + pub cluster_count: u16, + /// Embedding dimension after matryoshka truncation. Must be a positive + /// multiple of 8; values above 3072 are rejected. Defaults to 256. + #[serde(default = "ClusterEntitiesParams::default_dimension")] + pub dimension: u16, + + /// Seed for the random number generator used in clustering. + /// + /// If not provided, a random seed will be used. + pub seed: Option, +} + +impl ClusterEntitiesParams { + const fn default_dimension() -> u16 { + 256 + } +} + +/// One cluster from a spherical k-means run over entity embeddings. +#[derive(Debug, Serialize)] +#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))] +#[serde(rename_all = "camelCase")] +pub struct EntityCluster { + /// Index in `0..cluster_count`. + pub cluster_id: u16, + pub entity_ids: Vec, + /// Unit-normalized centroid with length equal to the requested dimension. + pub centroid: Vec, +} + +/// Result of [`EntityStore::cluster_entities`]. +#[derive(Debug, Serialize)] +#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))] +#[serde(rename_all = "camelCase")] +pub struct ClusterEntitiesResponse { + /// One entry per non-empty cluster. Empty clusters (no points assigned) + /// are omitted. + pub clusters: Vec, + /// Entities from the request that had no stored embedding. + pub missing_embeddings: Vec, +} + #[derive(Debug, Deserialize)] #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -912,6 +963,27 @@ pub trait EntityStore { params: UpdateEntityEmbeddingsParams<'_>, ) -> impl Future>> + Send; + /// Groups entities by embedding similarity using spherical k-means. + /// + /// Each entity's combined embedding is truncated to the requested + /// dimension (matryoshka encoding) before clustering. The returned + /// centroids are unit-normalized and have the same dimension. + /// + /// Entities without a stored embedding are not clustered; they appear + /// in [`ClusterEntitiesResponse::missing_embeddings`]. + /// + /// # Errors + /// + /// Returns [`ClusterError::InvalidDimension`] if the dimension is not a + /// positive multiple of 8, [`ClusterError::DimensionTooLarge`] if it + /// exceeds the stored embedding width, or [`ClusterError::Store`] if the + /// embedding query fails. + fn cluster_entities( + &self, + actor_id: ActorEntityUuid, + params: ClusterEntitiesParams, + ) -> impl Future>> + Send; + /// Re-indexes the cache for entities. /// /// This is only needed if the entity was changed in place without an update procedure. This is diff --git a/libs/@local/graph/store/src/error.rs b/libs/@local/graph/store/src/error.rs index d9f39229103..fc40ef8daf4 100644 --- a/libs/@local/graph/store/src/error.rs +++ b/libs/@local/graph/store/src/error.rs @@ -68,3 +68,18 @@ pub enum CheckPermissionError { } impl Error for CheckPermissionError {} + +/// Failure to cluster entities by embedding similarity. +#[derive(Debug, derive_more::Display)] +#[display("Could not cluster entities: {_variant}")] +#[must_use] +pub enum ClusterError { + #[display("dimension {dimension} is not a positive multiple of 8")] + InvalidDimension { dimension: u16 }, + #[display("dimension {dimension} exceeds stored embedding dimension {max}")] + DimensionTooLarge { dimension: u16, max: u16 }, + #[display("embedding query failed")] + Store, +} + +impl Error for ClusterError {} diff --git a/libs/@local/graph/store/src/lib.rs b/libs/@local/graph/store/src/lib.rs index 2931d3f0eca..668c46b20b8 100644 --- a/libs/@local/graph/store/src/lib.rs +++ b/libs/@local/graph/store/src/lib.rs @@ -5,6 +5,10 @@ #![feature( // Language Features impl_trait_in_assoc_type, + + // Library Features, + portable_simd, + integer_widen_truncate, )] #![cfg_attr(test, feature( // Language Features @@ -23,6 +27,7 @@ pub mod oauth_provider; pub mod property_type; pub mod user_deletion; +pub mod embedding; pub mod error; pub mod filter; pub mod migration; diff --git a/libs/@local/graph/type-fetcher/src/store.rs b/libs/@local/graph/type-fetcher/src/store.rs index 99a3ad3415c..d310af8f9a6 100644 --- a/libs/@local/graph/type-fetcher/src/store.rs +++ b/libs/@local/graph/type-fetcher/src/store.rs @@ -35,10 +35,10 @@ use hash_graph_store::{ UnarchiveDataTypeParams, UpdateDataTypeEmbeddingParams, UpdateDataTypesParams, }, entity::{ - CreateEntityParams, DeleteEntitiesParams, DeletionSummary, EntityStore, - EntityValidationReport, HasPermissionForEntitiesParams, PatchEntityParams, - QueryEntitiesParams, QueryEntitiesResponse, QueryEntitySubgraphParams, - QueryEntitySubgraphResponse, SearchEntitiesParams, SearchEntitiesResponse, + ClusterEntitiesParams, ClusterEntitiesResponse, CreateEntityParams, DeleteEntitiesParams, + DeletionSummary, EntityStore, EntityValidationReport, HasPermissionForEntitiesParams, + PatchEntityParams, QueryEntitiesParams, QueryEntitiesResponse, QueryEntitySubgraphParams, + QueryEntitySubgraphResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, ValidateEntityParams, }, @@ -50,7 +50,9 @@ use hash_graph_store::{ QueryEntityTypesResponse, SearchEntityTypesParams, SearchEntityTypesResponse, UnarchiveEntityTypeParams, UpdateEntityTypeEmbeddingParams, UpdateEntityTypesParams, }, - error::{CheckPermissionError, DeletionError, InsertionError, QueryError, UpdateError}, + error::{ + CheckPermissionError, ClusterError, DeletionError, InsertionError, QueryError, UpdateError, + }, filter::{Filter, QueryRecord}, pool::StorePool, property_type::{ @@ -1713,6 +1715,14 @@ where self.store.update_entity_embeddings(actor_id, params).await } + async fn cluster_entities( + &self, + actor_id: ActorEntityUuid, + params: ClusterEntitiesParams, + ) -> Result> { + self.store.cluster_entities(actor_id, params).await + } + async fn reindex_entity_cache(&mut self) -> Result<(), Report> { self.store.reindex_entity_cache().await } diff --git a/tests/graph/integration/postgres/lib.rs b/tests/graph/integration/postgres/lib.rs index 5c2ab899e3f..f9e54f423bb 100644 --- a/tests/graph/integration/postgres/lib.rs +++ b/tests/graph/integration/postgres/lib.rs @@ -890,6 +890,17 @@ impl EntityStore for DatabaseApi<'_> { self.store.reindex_entity_cache().await } + async fn cluster_entities( + &self, + actor_id: ActorEntityUuid, + params: hash_graph_store::entity::ClusterEntitiesParams, + ) -> Result< + hash_graph_store::entity::ClusterEntitiesResponse, + Report, + > { + self.store.cluster_entities(actor_id, params).await + } + async fn has_permission_for_entities( &self, authenticated_actor: AuthenticatedActor, From 4c7c562178b561f5ccd8ba555343393666b7d95d Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:07:07 +0200 Subject: [PATCH 02/38] fix: suggestions from code review --- Cargo.lock | 12 +++++++ Cargo.toml | 1 + .../store/postgres/knowledge/entity/mod.rs | 32 ++++++++----------- libs/@local/graph/store/Cargo.toml | 3 ++ .../graph/store/src/embedding/dimension.rs | 6 ++++ libs/@local/graph/store/src/entity/mod.rs | 3 +- libs/@local/graph/store/src/entity/store.rs | 7 ++-- libs/@local/graph/store/src/error.rs | 6 ++-- libs/@local/graph/type-fetcher/src/store.rs | 2 +- 9 files changed, 46 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c0d06b1b2d3..375c424b10a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3815,6 +3815,9 @@ dependencies = [ "hash-temporal-client", "insta", "postgres-types", + "rand 0.10.1", + "rand_xoshiro", + "rayon", "serde", "serde_json", "simple-mermaid", @@ -8032,6 +8035,15 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand_xoshiro" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "662effc7698e08ea324d3acccf8d9d7f7bf79b9785e270a174ea36e56900c91d" +dependencies = [ + "rand_core 0.10.1", +] + [[package]] name = "rapidfuzz" version = "0.5.0" diff --git a/Cargo.toml b/Cargo.toml index e97c0971742..49b388cbcf2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -224,6 +224,7 @@ quote = { version = "1.0.41", default-features = fa rand = { version = "0.10.0", default-features = false } rand_core = { version = "0.10.0", default-features = false } rand_distr = { version = "0.6.0", default-features = false } +rand_xoshiro = { version = "0.8.1" } rapidfuzz = { version = "0.5.0", default-features = false } ratatui = { version = "0.30.0" } rayon = { version = "1.11.0", default-features = false } diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index ff70b0a7d00..1f7b3b61f78 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -3,7 +3,7 @@ mod query; mod read; mod summary; -use alloc::borrow::Cow; +use alloc::{borrow::Cow, collections::BTreeMap}; use core::{borrow::Borrow as _, mem}; use std::collections::{HashMap, HashSet}; @@ -25,7 +25,8 @@ use hash_graph_store::{ EntityQueryPath, EntityQuerySorting, EntityStore, EntityTypeRetrieval, EntityTypesError, EntityValidationReport, EntityValidationType, HasPermissionForEntitiesParams, PatchEntityParams, QueryConversion, QueryEntitiesParams, QueryEntitiesResponse, - QueryEntitySubgraphParams, QueryEntitySubgraphResponse, SummarizeEntitiesParams, + QueryEntitySubgraphParams, QueryEntitySubgraphResponse, SearchEntitiesFilter, + SearchEntitiesParams, SearchEntitiesResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, ValidateEntityComponents, ValidateEntityParams, }, @@ -2598,36 +2599,31 @@ where Ok(permitted_ids) } - #[expect(clippy::too_many_lines)] + #[expect(clippy::too_many_lines, clippy::cast_possible_truncation)] #[tracing::instrument(skip(self, params))] async fn cluster_entities( &self, actor_id: ActorEntityUuid, params: ClusterEntitiesParams, ) -> Result> { - // 3072 fits in u16; compile-time verified. - const { - assert!(Embedding::DIM <= u16::MAX as usize); - } - #[expect( - clippy::cast_possible_truncation, - reason = "guarded by the const assertion above" - )] + const { assert!(Embedding::DIM <= u16::MAX as usize) }; const STORED_DIM: u16 = Embedding::DIM as u16; - let dim = Dimension::new(params.dimension).ok_or_else(|| { + let dimension = Dimension::new(params.dimension.get()).ok_or_else(|| { Report::new(ClusterError::InvalidDimension { dimension: params.dimension, }) + .attach(StatusCode::InvalidArgument) })?; - if dim.get() > STORED_DIM { + if dimension.get() > STORED_DIM { return Err(Report::new(ClusterError::DimensionTooLarge { - dimension: dim.get(), + dimension: dimension.value(), max: STORED_DIM, - })); + }) + .attach(StatusCode::InvalidArgument)); } - let truncated_dim = usize::from(dim.get()); + let truncated_dim = usize::from(dimension.get()); // Filter to entities the actor is allowed to view. let permitted = self @@ -2745,9 +2741,9 @@ where }), ); - let result = hash_graph_store::embedding::clustering::cluster(&flat, dim, &config); + let result = hash_graph_store::embedding::clustering::cluster(&flat, dimension, &config); - let mut groups: HashMap> = HashMap::new(); + let mut groups: BTreeMap> = BTreeMap::new(); for (index, id) in found_ids.iter().enumerate() { groups.entry(result.label(index)).or_default().push(*id); } diff --git a/libs/@local/graph/store/Cargo.toml b/libs/@local/graph/store/Cargo.toml index 825638c1e35..c02ac60c582 100644 --- a/libs/@local/graph/store/Cargo.toml +++ b/libs/@local/graph/store/Cargo.toml @@ -29,6 +29,9 @@ bytes = { workspace = true, optional = true } derive-where = { workspace = true } derive_more = { workspace = true, features = ["display", "error"] } futures = { workspace = true } +rand = { workspace = true } +rand_xoshiro = { workspace = true } +rayon = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } simple-mermaid = { workspace = true } diff --git a/libs/@local/graph/store/src/embedding/dimension.rs b/libs/@local/graph/store/src/embedding/dimension.rs index a3a1a31fe94..0ba85516ee3 100644 --- a/libs/@local/graph/store/src/embedding/dimension.rs +++ b/libs/@local/graph/store/src/embedding/dimension.rs @@ -31,6 +31,12 @@ impl Dimension { pub const fn get(self) -> u16 { self.0.get() } + + /// The raw dimension value as a [`NonZero`]. + #[must_use] + pub const fn value(self) -> NonZero { + self.0 + } } pub const D128: Dimension = Dimension(NonZero::new(128).unwrap()); diff --git a/libs/@local/graph/store/src/entity/mod.rs b/libs/@local/graph/store/src/entity/mod.rs index df563a365cc..e9aa8161500 100644 --- a/libs/@local/graph/store/src/entity/mod.rs +++ b/libs/@local/graph/store/src/entity/mod.rs @@ -9,7 +9,8 @@ pub use self::{ DiffEntityResult, EntityCluster, EntityPermissions, EntityStore, EntityValidationType, HasPermissionForEntitiesParams, LinkDeletionBehavior, PatchEntityParams, QueryConversion, QueryEntitiesParams, QueryEntitiesResponse, QueryEntitySubgraphParams, - QueryEntitySubgraphResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, + QueryEntitySubgraphResponse, SearchEntitiesFilter, SearchEntitiesParams, + SearchEntitiesResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, ValidateEntityComponents, ValidateEntityError, ValidateEntityParams, }, diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 1af2a5af1c4..91bccfdf907 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -1,4 +1,5 @@ use alloc::borrow::Cow; +use core::num::NonZero; use std::collections::{HashMap, HashSet}; use error_stack::Report; @@ -538,7 +539,7 @@ pub struct ClusterEntitiesParams { /// Embedding dimension after matryoshka truncation. Must be a positive /// multiple of 8; values above 3072 are rejected. Defaults to 256. #[serde(default = "ClusterEntitiesParams::default_dimension")] - pub dimension: u16, + pub dimension: NonZero, /// Seed for the random number generator used in clustering. /// @@ -547,8 +548,8 @@ pub struct ClusterEntitiesParams { } impl ClusterEntitiesParams { - const fn default_dimension() -> u16 { - 256 + const fn default_dimension() -> NonZero { + const { NonZero::new(256).unwrap() } } } diff --git a/libs/@local/graph/store/src/error.rs b/libs/@local/graph/store/src/error.rs index fc40ef8daf4..3f1b684167e 100644 --- a/libs/@local/graph/store/src/error.rs +++ b/libs/@local/graph/store/src/error.rs @@ -1,4 +1,4 @@ -use core::{error::Error, fmt}; +use core::{error::Error, fmt, num::NonZero}; #[derive(Debug)] #[must_use] @@ -75,9 +75,9 @@ impl Error for CheckPermissionError {} #[must_use] pub enum ClusterError { #[display("dimension {dimension} is not a positive multiple of 8")] - InvalidDimension { dimension: u16 }, + InvalidDimension { dimension: NonZero }, #[display("dimension {dimension} exceeds stored embedding dimension {max}")] - DimensionTooLarge { dimension: u16, max: u16 }, + DimensionTooLarge { dimension: NonZero, max: u16 }, #[display("embedding query failed")] Store, } diff --git a/libs/@local/graph/type-fetcher/src/store.rs b/libs/@local/graph/type-fetcher/src/store.rs index d310af8f9a6..87a124b74b0 100644 --- a/libs/@local/graph/type-fetcher/src/store.rs +++ b/libs/@local/graph/type-fetcher/src/store.rs @@ -38,7 +38,7 @@ use hash_graph_store::{ ClusterEntitiesParams, ClusterEntitiesResponse, CreateEntityParams, DeleteEntitiesParams, DeletionSummary, EntityStore, EntityValidationReport, HasPermissionForEntitiesParams, PatchEntityParams, QueryEntitiesParams, QueryEntitiesResponse, QueryEntitySubgraphParams, - QueryEntitySubgraphResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, + QueryEntitySubgraphResponse, SearchEntitiesParams, SearchEntitiesResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams, ValidateEntityParams, }, From c2c072824b7768e2ac609d997081656f6de41845 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:13:16 +0200 Subject: [PATCH 03/38] fix: spawn blocking for clustering --- .../src/store/postgres/knowledge/entity/mod.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 1f7b3b61f78..e850d1e14f6 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -2741,7 +2741,11 @@ where }), ); - let result = hash_graph_store::embedding::clustering::cluster(&flat, dimension, &config); + let result = tokio::task::spawn_blocking(move || { + hash_graph_store::embedding::clustering::cluster(&flat, dimension, &config) + }) + .await + .change_context(ClusterError::Store)?; let mut groups: BTreeMap> = BTreeMap::new(); for (index, id) in found_ids.iter().enumerate() { From a63a05a0cf811cac82596fc9073f65cc8438c0ea Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:14:49 +0200 Subject: [PATCH 04/38] fix: regenerate --- libs/@local/graph/api/openapi/openapi.json | 135 +++++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index 07d1ae085ab..1f224711a82 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -1580,6 +1580,54 @@ } } }, + "/entities/embeddings/clusters": { + "post": { + "tags": [ + "Graph", + "Entity" + ], + "operationId": "cluster_entities", + "parameters": [ + { + "name": "X-Authenticated-User-Actor-Id", + "in": "header", + "description": "The ID of the actor which is used to authorize the request", + "required": true, + "schema": { + "$ref": "#/components/schemas/ActorEntityUuid" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ClusterEntitiesParams" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Clusters of entities by embedding similarity", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ClusterEntitiesResponse" + } + } + } + }, + "422": { + "description": "Provided request body is invalid" + }, + "500": { + "description": "Store error occurred" + } + } + } + }, "/entities/permissions": { "post": { "tags": [ @@ -3742,6 +3790,62 @@ "propertyName": "kind" } }, + "ClusterEntitiesParams": { + "type": "object", + "required": [ + "entityIds", + "clusterCount" + ], + "properties": { + "clusterCount": { + "type": "integer", + "format": "int32", + "description": "Desired number of clusters. Clamped to the number of entities with\nembeddings when that is smaller.", + "minimum": 0 + }, + "dimension": { + "$ref": "#/components/schemas/NonZero" + }, + "entityIds": { + "type": "array", + "items": { + "$ref": "#/components/schemas/EntityId" + } + }, + "seed": { + "type": "integer", + "format": "int64", + "description": "Seed for the random number generator used in clustering.\n\nIf not provided, a random seed will be used.", + "nullable": true, + "minimum": 0 + } + }, + "additionalProperties": false + }, + "ClusterEntitiesResponse": { + "type": "object", + "description": "Result of [`EntityStore::cluster_entities`].", + "required": [ + "clusters", + "missingEmbeddings" + ], + "properties": { + "clusters": { + "type": "array", + "items": { + "$ref": "#/components/schemas/EntityCluster" + }, + "description": "One entry per non-empty cluster. Empty clusters (no points assigned)\nare omitted." + }, + "missingEmbeddings": { + "type": "array", + "items": { + "$ref": "#/components/schemas/EntityId" + }, + "description": "Entities from the request that had no stored embedding." + } + } + }, "CommonQueryEntityTypesParams": { "type": "object", "required": [ @@ -4507,6 +4611,37 @@ } } }, + "EntityCluster": { + "type": "object", + "description": "One cluster from a spherical k-means run over entity embeddings.", + "required": [ + "clusterId", + "entityIds", + "centroid" + ], + "properties": { + "centroid": { + "type": "array", + "items": { + "type": "number", + "format": "float" + }, + "description": "Unit-normalized centroid with length equal to the requested dimension." + }, + "clusterId": { + "type": "integer", + "format": "int32", + "description": "Index in `0..cluster_count`.", + "minimum": 0 + }, + "entityIds": { + "type": "array", + "items": { + "$ref": "#/components/schemas/EntityId" + } + } + } + }, "EntityDeletionProvenance": { "type": "object", "required": [ From c1c0f7feda83491a0f9e30372ae445dfb7879a5c Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:22:47 +0200 Subject: [PATCH 05/38] fix: unnest sql query --- .../postgres-store/src/store/postgres/knowledge/entity/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index e850d1e14f6..50311503127 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -2665,8 +2665,8 @@ where embedding FROM entity_embeddings e WHERE e.property IS NULL - AND (e.web_id, e.entity_uuid) IN (SELECT unnest($1::uuid[]), \ - unnest($2::uuid[]))" + AND (e.web_id, e.entity_uuid) IN (SELECT * FROM unnest($1::uuid[], \ + $2::uuid[]))" ), [ &web_ids as &(dyn ToSql + Sync), From 3b5b26e10fe73f274fa6ea8a0aeb946cb34c7504 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:31:18 +0200 Subject: [PATCH 06/38] fix: lints --- libs/@local/graph/store/src/embedding/clustering.rs | 11 ----------- libs/@local/graph/store/src/embedding/kernel.rs | 11 ++++++----- libs/@local/graph/store/src/embedding/mod.rs | 6 ++---- 3 files changed, 8 insertions(+), 20 deletions(-) diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index f48da06398a..45839919a68 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -37,12 +37,6 @@ pub struct Config { } impl Config { - /// Creates a configuration for `k` clusters, drawing the seed from `rng`. - #[must_use] - pub(crate) fn for_k(k: u16, mut rng: impl Rng) -> Self { - Self::for_k_with_seed(k, rng.random()) - } - /// Creates a configuration for `k` clusters with a fixed seed. /// /// Defaults: 30 max iterations, 5 restarts, 1e-4 convergence tolerance, @@ -112,11 +106,6 @@ impl Clustering { pub fn label(&self, entity: usize) -> u16 { self.labels[entity] } - - /// Returns a mutable reference to the cluster label for point `entity`. - fn label_mut(&mut self, entity: usize) -> &mut u16 { - &mut self.labels[entity] - } } // TODO: I wonder if we can make this allocation less diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/store/src/embedding/kernel.rs index 0090f63c4c5..7abaf1c8b5b 100644 --- a/libs/@local/graph/store/src/embedding/kernel.rs +++ b/libs/@local/graph/store/src/embedding/kernel.rs @@ -1,3 +1,8 @@ +#![expect( + clippy::inline_always, + reason = "while usually discouraged, SIMD operations need to be inlined, as otherwise we \ + spill SIMD registers, see the SIMD documentation." +)] use core::simd::{Simd, f32x8, num::SimdFloat as _}; /// Fused multiply-add when the target has native FMA, separate mul+add otherwise. @@ -241,11 +246,7 @@ pub(crate) unsafe fn normalize(value: &mut [f32]) { /// # Safety /// * all six slices have length `d` /// * `d` is a multiple of 8 -#[expect( - clippy::inline_always, - reason = "micro-kernel must inline into nearest4 to keep accumulators in registers" -)] -#[inline(always)] +#[inline(always)] // micro-kernel must inline nearest4 to keep accumulators in registers pub(crate) unsafe fn micro_4x2( p0: &[f32], p1: &[f32], diff --git a/libs/@local/graph/store/src/embedding/mod.rs b/libs/@local/graph/store/src/embedding/mod.rs index cefb998fb67..d008b247ce0 100644 --- a/libs/@local/graph/store/src/embedding/mod.rs +++ b/libs/@local/graph/store/src/embedding/mod.rs @@ -1,13 +1,11 @@ #![expect( unsafe_code, - dead_code, clippy::indexing_slicing, clippy::float_arithmetic, clippy::min_ident_chars, clippy::many_single_char_names, - reason = "embedding module is under active development; dead_code is expected until the \ - public API is wired up. Single-char idents (k, n, m, d, x) are standard \ - mathematical notation for clustering." + reason = "embedding module is under active development. Single-char idents (k, n, m, d, x) \ + are standard mathematical notation for clustering." )] pub mod clustering; From d1c214ef5b772dc7edd6155c27fa5cd9be2caa41 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 11:02:03 +0200 Subject: [PATCH 07/38] fix: lints --- libs/@local/graph/store/src/embedding/kernel.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/store/src/embedding/kernel.rs index 7abaf1c8b5b..d5610052694 100644 --- a/libs/@local/graph/store/src/embedding/kernel.rs +++ b/libs/@local/graph/store/src/embedding/kernel.rs @@ -7,8 +7,8 @@ use core::simd::{Simd, f32x8, num::SimdFloat as _}; /// Fused multiply-add when the target has native FMA, separate mul+add otherwise. /// -/// On aarch64, FMA is part of the base NEON instruction set (`fmla`). -/// On x86_64, FMA requires the `fma` target feature (`vfmadd`); without it, +/// On `aarch64`, FMA is part of the base NEON instruction set (`fmla`). +/// On `x86_64`, FMA requires the `fma` target feature (`vfmadd`); without it, /// `StdFloat::mul_add` falls back to a per-lane `fmaf` libc call which /// destroys throughput. The non-FMA path uses a plain multiply and add /// (`vmulps` + `vaddps`) instead. From 2249f1bd9293e0725e03535d1dbddc72bba0aa12 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 1 Jul 2026 11:06:51 +0200 Subject: [PATCH 08/38] fix: openapi schema --- libs/@local/graph/api/openapi/openapi.json | 7 ++++++- libs/@local/graph/store/src/entity/store.rs | 1 + 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index 1f224711a82..365193087c7 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -3804,7 +3804,12 @@ "minimum": 0 }, "dimension": { - "$ref": "#/components/schemas/NonZero" + "type": "integer", + "format": "int32", + "description": "Embedding dimension after matryoshka truncation. Must be a positive\nmultiple of 8; values above 3072 are rejected. Defaults to 256.", + "default": 256, + "example": 256, + "minimum": 1 }, "entityIds": { "type": "array", diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 91bccfdf907..e6c43cdea12 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -539,6 +539,7 @@ pub struct ClusterEntitiesParams { /// Embedding dimension after matryoshka truncation. Must be a positive /// multiple of 8; values above 3072 are rejected. Defaults to 256. #[serde(default = "ClusterEntitiesParams::default_dimension")] + #[cfg_attr(feature = "utoipa", schema(value_type = u16, minimum = 1, default = 256, example = 256))] pub dimension: NonZero, /// Seed for the random number generator used in clustering. From 9b7fef660c81c2e762656de95060d882af7ee928 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Fri, 3 Jul 2026 13:09:22 +0200 Subject: [PATCH 09/38] fix: docs --- libs/@local/graph/store/src/embedding/clustering.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index 45839919a68..cec24ed18ab 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -9,8 +9,8 @@ use super::{dimension::Dimension, kernel}; /// Parameters for k-means clustering. /// -/// Use [`Config::for_k`] or [`Config::for_k_with_seed`] to construct with -/// reasonable defaults, then override individual fields as needed. +/// Use [`Config::for_k_with_seed`] to construct with reasonable defaults, then override individual +/// fields as needed. pub struct Config { /// Number of clusters. pub k: u16, From 121cd2293669338dc77ca9ba17a49d61a1e7abf2 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Fri, 3 Jul 2026 13:11:31 +0200 Subject: [PATCH 10/38] fix: docs --- libs/@local/graph/store/src/embedding/kernel.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/store/src/embedding/kernel.rs index d5610052694..1b971bc23b1 100644 --- a/libs/@local/graph/store/src/embedding/kernel.rs +++ b/libs/@local/graph/store/src/embedding/kernel.rs @@ -320,9 +320,8 @@ pub(crate) unsafe fn micro_4x2( /// [`micro_4x2`] tiled kernel. /// /// Returns `(centroid_index, raw_dot_product)` for each of the 4 points. -/// The raw dot product is **not** a distance; the caller must convert via -/// [`squared_chord_distance`](super::clustering::squared_chord_distance) -/// if needed. +/// The raw dot product is **not** a distance; and must be converted via +/// using the chord distance formula. /// /// # Safety /// From 9bffac9cefeef55860e908fbe30748888d3cd87d Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Fri, 3 Jul 2026 16:21:06 +0200 Subject: [PATCH 11/38] feat: embedding clustering review --- Cargo.lock | 2 + libs/@local/graph/api/openapi/openapi.json | 8 +- .../store/postgres/knowledge/entity/mod.rs | 6 + libs/@local/graph/store/Cargo.toml | 12 +- libs/@local/graph/store/benches/embedding.rs | 237 ++++ .../graph/store/src/embedding/clustering.rs | 1260 ++++++++++------- .../graph/store/src/embedding/kernel.rs | 279 ++-- libs/@local/graph/store/src/embedding/mod.rs | 5 +- libs/@local/graph/store/src/entity/store.rs | 5 + 9 files changed, 1145 insertions(+), 669 deletions(-) create mode 100644 libs/@local/graph/store/benches/embedding.rs diff --git a/Cargo.lock b/Cargo.lock index 375c424b10a..8c32c3395bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3803,6 +3803,8 @@ name = "hash-graph-store" version = "0.0.0" dependencies = [ "bytes", + "codspeed-criterion-compat", + "darwin-kperf-criterion", "derive-where", "derive_more", "error-stack", diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index 365193087c7..a93edfd1e59 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -3832,7 +3832,8 @@ "description": "Result of [`EntityStore::cluster_entities`].", "required": [ "clusters", - "missingEmbeddings" + "missingEmbeddings", + "inertia" ], "properties": { "clusters": { @@ -3842,6 +3843,11 @@ }, "description": "One entry per non-empty cluster. Empty clusters (no points assigned)\nare omitted." }, + "inertia": { + "type": "number", + "format": "float", + "description": "Sum of squared chord distances from every clustered entity to its\nassigned centroid. Lower is tighter; comparable across runs over the\nsame entities, e.g. to choose a cluster count. `0.0` when nothing was\nclustered." + }, "missingEmbeddings": { "type": "array", "items": { diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 50311503127..4faecf6befb 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -2654,6 +2654,10 @@ where // Truncate server-side via `subvector` so postgres only sends // `truncated_dim`-dimensional vectors over the wire. + // + // Matryoshka truncation shortens the vectors without re-normalizing; + // that is fine here because spherical k-means normalizes internally + // (it works with inverse norms), so no `l2_normalize` is needed. let row_stream = self .as_client() .query_raw( @@ -2722,6 +2726,7 @@ where return Ok(ClusterEntitiesResponse { clusters: Vec::new(), missing_embeddings, + inertia: 0.0, }); } @@ -2764,6 +2769,7 @@ where Ok(ClusterEntitiesResponse { clusters, missing_embeddings, + inertia: result.inertia, }) } } diff --git a/libs/@local/graph/store/Cargo.toml b/libs/@local/graph/store/Cargo.toml index c02ac60c582..afd99bc632c 100644 --- a/libs/@local/graph/store/Cargo.toml +++ b/libs/@local/graph/store/Cargo.toml @@ -40,14 +40,20 @@ tracing = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] -hash-codegen = { workspace = true } -insta = { workspace = true } -tokio = { workspace = true, features = ["macros"] } +codspeed-criterion-compat = { workspace = true } +darwin-kperf-criterion = { workspace = true, features = ["codspeed"] } +hash-codegen = { workspace = true } +insta = { workspace = true } +tokio = { workspace = true, features = ["macros"] } [[test]] name = "codegen" required-features = ["codegen"] +[[bench]] +name = "embedding" +harness = false + [features] codegen = ["dep:specta", "type-system/codegen", "hash-graph-authorization/codegen"] utoipa = ["hash-graph-temporal-versioning/utoipa", "type-system/utoipa", "dep:utoipa"] diff --git a/libs/@local/graph/store/benches/embedding.rs b/libs/@local/graph/store/benches/embedding.rs new file mode 100644 index 00000000000..13094b45589 --- /dev/null +++ b/libs/@local/graph/store/benches/embedding.rs @@ -0,0 +1,237 @@ +//! Benchmarks for the embedding k-means module. +//! +//! Two groups: +//! +//! * `embedding/kernel/*` — single-threaded SIMD micro-kernels, measured in retired instructions +//! via Apple PMCs (near-deterministic; requires root on macOS) with an automatic wall-clock +//! fallback on other platforms. +//! * `embedding/cluster/*` — end-to-end [`cluster`] runs. Always wall-clock, because the work is +//! spread across the rayon pool and per-thread instruction counts would only see the calling +//! thread. +//! +//! [`cluster`]: hash_graph_store::embedding::clustering::cluster +#![expect( + unsafe_code, + clippy::float_arithmetic, + clippy::indexing_slicing, + clippy::integer_division, + clippy::integer_division_remainder_used, + clippy::min_ident_chars, + clippy::significant_drop_tightening, + reason = "benchmarks exercise the unsafe SIMD kernels directly and build float test data; \ + single-char idents (k, n, d) are standard mathematical notation for clustering; the \ + drop-tightening warning originates inside `criterion_group!`" +)] + +use core::hint::black_box; + +use codspeed_criterion_compat::{ + BenchmarkId, Criterion, criterion_group, criterion_main, measurement::Measurement, +}; +use hash_graph_store::embedding::{ + clustering::{Config, cluster}, + dimension::Dimension, + kernel, +}; +use rand::{RngExt as _, SeedableRng as _}; +use rand_xoshiro::Xoshiro256PlusPlus; + +/// Uniform random values in `[-1, 1)`. +fn random_vec(len: usize, seed: u64) -> Vec { + let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); + core::iter::repeat_with(|| rng.random_range(-1.0..1.0)) + .take(len) + .collect() +} + +/// Uniform random values in `[0.1, 1)`, guaranteed positive so repeated +/// accumulation saturates at infinity instead of producing NaNs. +fn random_positive_vec(len: usize, seed: u64) -> Vec { + let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); + core::iter::repeat_with(|| rng.random_range(0.1..1.0)) + .take(len) + .collect() +} + +/// Well-separated blobs: `k` clusters of `points_per_cluster` points in +/// `d`-dimensional space, each with a dominant axis. Mirrors the shape of +/// real embedding workloads better than uniform noise: the fit converges +/// instead of always exhausting `max_iters`. +fn blobs(points_per_cluster: usize, k: usize, d: usize, seed: u64) -> Vec { + let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); + let mut data = vec![0.0_f32; points_per_cluster * k * d]; + + for (index, row) in data.chunks_exact_mut(d).enumerate() { + let axis = (index / points_per_cluster) % d; + row[axis] = 10.0; + for value in row.iter_mut() { + *value += rng.random_range(-0.01..0.01); + } + } + + data +} + +const KERNEL_DIMS: &[usize] = &[256, 1536, 3072]; + +fn bench_dot(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("embedding/kernel/dot"); + + for &d in KERNEL_DIMS { + let lhs = random_vec(d, 1); + let rhs = random_vec(d, 2); + + group.bench_with_input(BenchmarkId::from_parameter(d), &d, |bencher, _| { + // SAFETY: both slices have length `d`, a multiple of 8. + bencher.iter(|| unsafe { kernel::dot(black_box(&lhs), black_box(&rhs)) }); + }); + } + + group.finish(); +} + +fn bench_add_scaled_into(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("embedding/kernel/add_scaled_into"); + + for &d in KERNEL_DIMS { + let src = random_positive_vec(d, 3); + let mut dst = random_positive_vec(d, 4); + + group.bench_with_input(BenchmarkId::from_parameter(d), &d, |bencher, _| { + // SAFETY: both slices have length `d`, a multiple of 8. + bencher.iter(|| unsafe { + kernel::add_scaled_into(black_box(&mut dst), black_box(&src), black_box(0.5)); + }); + }); + } + + group.finish(); +} + +fn bench_micro_4x2(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("embedding/kernel/micro_4x2"); + + for &d in KERNEL_DIMS { + let points: Vec> = (0..4).map(|seed| random_vec(d, 10 + seed)).collect(); + let c0 = random_vec(d, 20); + let c1 = random_vec(d, 21); + + group.bench_with_input(BenchmarkId::from_parameter(d), &d, |bencher, _| { + // SAFETY: all six slices have length `d`, a multiple of 8. + bencher.iter(|| unsafe { + kernel::micro_4x2( + black_box(&points[0]), + black_box(&points[1]), + black_box(&points[2]), + black_box(&points[3]), + black_box(&c0), + black_box(&c1), + ) + }); + }); + } + + group.finish(); +} + +macro_rules! nz { + ($expr:expr) => { + const { ::core::num::NonZero::new($expr).unwrap() } + }; +} + +fn bench_nearest4(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("embedding/kernel/nearest4"); + + // k = 15 exercises the odd-k remainder path. + for &(d, k) in &[ + (256, nz!(15)), + (256, nz!(16)), + (256, nz!(64)), + (1536, nz!(16)), + (3072, nz!(16)), + ] { + let points: Vec> = (0..4).map(|seed| random_vec(d, 30 + seed)).collect(); + let centroids = random_vec(k.get() * d, 40); + + group.bench_with_input( + BenchmarkId::new(format!("d{d}"), k), + &(d, k), + |bencher, _| { + // SAFETY: point slices have length `d` (multiple of 8), + // centroids has length `k * d`, and `k > 0`. + bencher.iter(|| unsafe { + kernel::nearest4( + black_box(&points[0]), + black_box(&points[1]), + black_box(&points[2]), + black_box(&points[3]), + black_box(¢roids), + black_box(k), + black_box(d), + ) + }); + }, + ); + } + + group.finish(); +} + +fn bench_cluster(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("embedding/cluster"); + group.sample_size(10); + + let dimension = Dimension::new(256).expect("256 is a positive multiple of 8"); + + // (n, k): n = 10k exercises the subsampled fit (m = 8192) plus the + // full-data refinement; n = 50k shifts the weight onto the full-data + // passes. + for &(n, k) in &[ + (10_000_usize, 8_u16), + (10_000, 32), + (10_000, 128), + (50_000, 32), + ] { + let data = blobs(n / usize::from(k), usize::from(k), 256, 7); + let config = Config::for_k_with_seed(k, 42); + + group.bench_with_input( + BenchmarkId::new(format!("n{n}_d256"), k), + &(n, k), + |bencher, _| { + bencher.iter(|| cluster(black_box(&data), black_box(dimension), &config)); + }, + ); + } + + group.finish(); +} + +fn kernel_measurement() -> Criterion { + use core::time::Duration; + + // Retired instructions on Apple Silicon (needs root there), wall-clock + // fallback everywhere else. Instruction counts are near-deterministic, + // so short windows and small samples suffice. + Criterion::default() + .with_measurement( + darwin_kperf_criterion::HardwareCounter::instructions() + .expect("instruction counting requires root on Apple Silicon (run under sudo)"), + ) + .warm_up_time(Duration::from_millis(500)) + .measurement_time(Duration::from_secs(1)) + .sample_size(20) +} + +criterion_group!( + name = kernel; + config = kernel_measurement(); + targets = bench_dot, bench_add_scaled_into, bench_micro_4x2, bench_nearest4 +); +criterion_group!( + name = clustering; + config = Criterion::default(); + targets = bench_cluster +); +criterion_main!(kernel, clustering); diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index cec24ed18ab..0ff02308ee7 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -1,9 +1,16 @@ use alloc::borrow::Cow; -use core::{cmp, mem, num::NonZero}; +use core::{cmp, num::NonZero}; +use std::collections::HashSet; use rand::{Rng, RngExt as _, SeedableRng as _}; use rand_xoshiro::Xoshiro256PlusPlus; -use rayon::prelude::*; +use rayon::{ + iter::{ + IndexedParallelIterator as _, IntoParallelIterator as _, IntoParallelRefIterator as _, + IntoParallelRefMutIterator as _, ParallelIterator as _, + }, + slice::{ParallelSlice as _, ParallelSliceMut as _}, +}; use super::{dimension::Dimension, kernel}; @@ -29,10 +36,14 @@ pub struct Config { /// Capped to avoid quadratic seeding cost on very large datasets. pub sample_cap: usize, - /// Base seed for the PRNG. Each restart derives its own seed from this. + /// Base seed for the PRNG. + /// + /// Runs with the same seed, input, and configuration produce identical + /// labels and centroids. pub seed: u64, - /// Number of points processed per batch in the assignment step. + /// Number of points processed per batch in the parallel passes. + /// Values larger than the number of points are clamped. pub chunk: NonZero, } @@ -64,10 +75,26 @@ pub struct Clustering { pub dimension: Dimension, /// Flat centroid matrix, `k * d` elements in row-major order. + /// + /// Centroids are unit-normalized, with one exception: a cluster whose + /// members are all zero-norm points keeps a zero centroid, since there + /// is no direction to normalize. pub centroids: Box<[f32]>, /// Cluster assignment for each input point, values in `0..k`. + /// + /// When [`cluster`] ran with `k == 0` (requested or clamped) the labels + /// are all-zero placeholders and there are no centroids to index. pub labels: Box<[u16]>, + + /// Sum of squared chord distances from every input point to its assigned + /// centroid, measured against the final centroids. Lower is tighter; + /// comparable across runs on the same input, e.g. for choosing `k`. + /// `0.0` when `k == 0` or the input is empty. + /// + /// The value is precise only up to floating-point summation order: + /// repeated runs over identical input can differ in the final bits. + pub inertia: f32, } impl Clustering { @@ -85,10 +112,15 @@ impl Clustering { centroids, labels, dimension: d, + inertia: 0.0, } } /// Returns the `D`-dimensional slice for centroid `cluster`. + /// + /// # Panics + /// + /// Panics if `cluster` is not below the number of centroids. #[must_use] pub fn centroid(&self, cluster: u16) -> &[f32] { &self.centroids[cluster as usize * (self.dimension.get() as usize) @@ -102,23 +134,40 @@ impl Clustering { } /// Returns the cluster label for point `entity`. + /// + /// # Panics + /// + /// Panics if `entity` is not below the number of input points. #[must_use] pub fn label(&self, entity: usize) -> u16 { self.labels[entity] } } -// TODO: I wonder if we can make this allocation less +/// Draws `m` distinct indices uniformly at random from `0..n` in O(m) time +/// and memory (Robert Floyd's sampling algorithm). +/// +/// The result is sorted and deterministic for a given RNG state. fn sample_indices(n: usize, m: usize, mut rng: impl Rng) -> Vec { - let mut idx: Vec = (0..n).collect(); + debug_assert!(m <= n); + + let mut selected: HashSet = HashSet::with_capacity(m); + + for upper in n - m..n { + let candidate = rng.random_range(0..=upper); - for i in 0..m { - let j = i + rng.random_range(0..n - i); // partial Fisher–Yates - idx.swap(i, j); + if !selected.insert(candidate) { + // `candidate` was already drawn in an earlier round. Earlier + // rounds only drew from `0..upper`, so `upper` itself is fresh. + selected.insert(upper); + } } - idx.truncate(m); - idx + let mut indices: Vec = selected.into_iter().collect(); + // Sorting erases the hash set's nondeterministic iteration order and + // turns the caller's gather into a forward walk over `x`. + indices.sort_unstable(); + indices } /// Squared chord distance between a point and a unit centroid. @@ -147,36 +196,33 @@ fn squared_chord_distance(dot: f32, point_inv_norm: f32) -> f32 { /// /// # Safety /// -/// * `point.len() == D` -/// * `centroids.len() == k * D` -/// * `k > 0` -/// * `D` is a multiple of 8 (enforced at compile time by the const generic). +/// * `point.len() == d` +/// * `centroids.len() == k * d` +/// * `d` is a multiple of 8 (guaranteed by [`Dimension`]). #[inline] #[must_use] pub(crate) unsafe fn nearest_centroid( point: &[f32], point_inv_norm: f32, centroids: &[f32], - k: usize, + k: NonZero, d: usize, ) -> (u16, f32) { debug_assert_eq!(point.len(), d); - debug_assert_eq!(centroids.len(), k * d); - debug_assert!(k > 0); + debug_assert_eq!(centroids.len(), k.get() * d); // SAFETY: the caller guarantees these preconditions. The hints let the // compiler elide bounds checks on the centroid slicing inside the loop. unsafe { core::hint::assert_unchecked(point.len() == d); - core::hint::assert_unchecked(centroids.len() == k * d); + core::hint::assert_unchecked(centroids.len() == k.get() * d); core::hint::assert_unchecked(d.is_multiple_of(8)); - core::hint::assert_unchecked(k > 0); } let mut best = 0; let mut best_dot = f32::NEG_INFINITY; - for cluster in 0..k { + for cluster in 0..k.get() { let start = cluster * d; let centroid = ¢roids[start..start + d]; @@ -186,7 +232,7 @@ pub(crate) unsafe fn nearest_centroid( #[expect( clippy::cast_possible_truncation, - reason = "k is supposed to be low, and checked as such via the config" + reason = "cluster < k, and k originates from Config::k (u16)" )] if dot > best_dot { best = cluster as u16; @@ -197,351 +243,182 @@ pub(crate) unsafe fn nearest_centroid( (best, squared_chord_distance(best_dot, point_inv_norm)) } -/// Pre-allocated scratch space for the k-means fitting loop. +/// Assigns one chunk of points during Lloyd iterations: writes each point's +/// nearest centroid into `labels` and its squared chord distance into +/// `distances`. +/// +/// # Safety +/// +/// * `points.len() == labels.len() * d` +/// * `inv_norms.len() == labels.len()` +/// * `distances.len() == labels.len()` +/// * `centroids.len() == k * d` +/// * `d` is a multiple of 8 +unsafe fn lloyd_assign( + k: NonZero, + d: usize, + centroids: &[f32], + points: &[f32], + inv_norms: &[f32], + labels: &mut [u16], + distances: &mut [f32], +) { + let count = labels.len(); + + // SAFETY: the caller guarantees the length relations; the hints let the + // compiler elide bounds checks in the tiled loop below. + unsafe { + core::hint::assert_unchecked(points.len() == count * d); + core::hint::assert_unchecked(inv_norms.len() == count); + core::hint::assert_unchecked(distances.len() == count); + core::hint::assert_unchecked(d.is_multiple_of(8)); + } + + let mut i = 0; + while i + 4 <= count { + let p0 = &points[i * d..i * d + d]; + let p1 = &points[(i + 1) * d..(i + 1) * d + d]; + let p2 = &points[(i + 2) * d..(i + 2) * d + d]; + let p3 = &points[(i + 3) * d..(i + 3) * d + d]; + + // SAFETY: each point length d, centroids length k*d, + // k > 0, d a multiple of 8 (guaranteed by Dimension). + let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; + + for offset in 0..4 { + labels[i + offset] = nearest[offset].0; + distances[i + offset] = + squared_chord_distance(nearest[offset].1, inv_norms[i + offset]); + } + i += 4; + } + + while i < count { + let point = &points[i * d..i * d + d]; + + // SAFETY: point length d, centroids length k*d, k > 0, d mult of 8. + let (label, distance) = unsafe { nearest_centroid(point, inv_norms[i], centroids, k, d) }; + labels[i] = label; + distances[i] = distance; + i += 1; + } +} + +/// Per-restart scratch and state for one k-means fit on the sample. /// -/// All buffers are allocated once and reused across restarts to avoid -/// per-iteration allocation overhead. -struct Fit { - k: usize, +/// Restarts run in parallel, so each owns its buffers. [`Restart::new`] +/// hands them back zeroed. +struct Restart { + k: NonZero, m: usize, d: usize, - /// Current centroids for this restart, `k * d` elements. + /// Centroids for this restart, `k * d` elements. centroids: Box<[f32]>, - /// Best centroids seen across all restarts. - best_centroids: Box<[f32]>, /// Per-cluster accumulator for centroid recomputation, `k * d` elements. sums: Box<[f32]>, - /// Per-cluster point count for centroid averaging. + /// Per-cluster point count for the empty-cluster check. counts: Box<[usize]>, /// Per-sample-point cluster assignment. labels: Box<[u16]>, - /// Per-sample-point closest centroid distance (for k-means++ seeding). - closest_distances: Box<[f32]>, + /// Per-sample-point distance scratch. + point_distances: Box<[f32]>, /// Tracks which sample points have been selected as seeds. selected: Box<[bool]>, - /// Lowest inertia across all restarts. - best_inertia: f32, } -impl Fit { - fn new(k: usize, m: usize, d: usize) -> Self { - // SAFETY: all-zero bits are valid for f32 (IEEE 754 +0.0), usize (0), u16 (0), and bool - // (false). `Box::new_zeroed_slice` allocates zeroed memory of the correct layout - // for each type, so `assume_init` is sound in every case. - let centroids = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; - // SAFETY: see above - let best_centroids = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; +impl Restart { + fn new(k: NonZero, m: usize, d: usize) -> Self { + // SAFETY: all-zero bits are valid for `f32` (IEEE 754 +0.0), `usize` (0), `u16` (0), and + // `bool` (false). `Box::new_zeroed_slice` allocates zeroed memory of the correct + // layout for each type, so `assume_init` is sound in every case. + let centroids = unsafe { Box::<[f32]>::new_zeroed_slice(k.get() * d).assume_init() }; // SAFETY: see above - let sums = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; + let sums = unsafe { Box::<[f32]>::new_zeroed_slice(k.get() * d).assume_init() }; // SAFETY: see above - let counts = unsafe { Box::<[usize]>::new_zeroed_slice(k).assume_init() }; + let counts = unsafe { Box::<[usize]>::new_zeroed_slice(k.get()).assume_init() }; // SAFETY: see above let labels = unsafe { Box::<[u16]>::new_zeroed_slice(m).assume_init() }; // SAFETY: see above - let closest_distances = unsafe { Box::<[f32]>::new_zeroed_slice(m).assume_init() }; + let point_distances = unsafe { Box::<[f32]>::new_zeroed_slice(m).assume_init() }; // SAFETY: see above let selected = unsafe { Box::<[bool]>::new_zeroed_slice(m).assume_init() }; - let best_inertia = f32::INFINITY; Self { k, m, d, centroids, - best_centroids, sums, counts, labels, - closest_distances, + point_distances, selected, - best_inertia, } } - fn reset_centroids(&mut self) { - self.centroids.fill(0.0); - } - - fn reset_sums(&mut self) { - self.sums.fill(0.0); - } - - fn reset_counts(&mut self) { - self.counts.fill(0); - } - - fn reset_selected(&mut self) { - self.selected.fill(false); - } - - /// Reinitializes empty clusters from the sample point farthest from - /// its assigned centroid. + /// Runs one restart: k-means++ seeding followed by Lloyd iterations. /// - /// For each empty cluster, scans the sample to find the point with - /// the largest squared chord distance to its current centroid, copies - /// that point as the new centroid (normalized), and updates the - /// point's label so it won't be picked again for subsequent empty - /// clusters in the same pass. - #[expect( - clippy::cast_possible_truncation, - reason = "cluster index < k, and k originates from Config::k (u16)" - )] - fn reinit_empty_clusters(&mut self, sample: &[f32], sample_inv_norms: &[f32]) -> bool { - let &mut Self { d, k, .. } = self; - - let mut reseeded = false; - - for cluster in 0..k { - if self.counts[cluster] != 0 { - continue; - } - - reseeded = true; - let mut farthest_idx = 0; - let mut farthest_dist = -1.0_f32; - - for (i, (point, &inv_norm)) in sample.chunks_exact(d).zip(sample_inv_norms).enumerate() - { - let label = usize::from(self.labels[i]); - let c_start = label * d; - - // SAFETY: point and centroid both have length `d`, - // a multiple of 8 (guaranteed by Dimension). - let dot = unsafe { kernel::dot(point, &self.centroids[c_start..c_start + d]) }; - let dist = squared_chord_distance(dot, inv_norm); - - if dist > farthest_dist { - farthest_dist = dist; - farthest_idx = i; - } - } - - let point_start = farthest_idx * d; - let centroid_start = cluster * d; - self.centroids[centroid_start..centroid_start + d] - .copy_from_slice(&sample[point_start..point_start + d]); - - // SAFETY: centroid row has length `d`, a multiple of 8. - unsafe { - kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); - } - - // Update the label so the next empty cluster picks a different - // point (this point's distance to its new centroid is ~0). - self.labels[farthest_idx] = cluster as u16; - } - - reseeded - } - - /// Runs k-means++ initialization followed by Lloyd iterations on the - /// sample, repeating for `n_init` restarts. The best centroids (lowest - /// inertia) are stored in `self.best_centroids`. + /// Returns the sample inertia of the fitted centroids. fn run( &mut self, sample: &[f32], chunk: usize, row_chunk: usize, sample_inv_norms: &[f32], - mut rng: impl Rng, - config: &Config, - ) { - for _ in 0..config.n_init.get() { - self.reset_centroids(); - self.closest_distances.fill(f32::INFINITY); - self.reset_selected(); - - self.seed_plusplus(sample, sample_inv_norms, &mut rng); - - let inertia = self.lloyd(sample, chunk, row_chunk, sample_inv_norms, config); - - if inertia < self.best_inertia { - self.best_inertia = inertia; - mem::swap(&mut self.best_centroids, &mut self.centroids); - } - } - } - - /// Runs Lloyd iterations on the sample until convergence or `max_iters`. - /// Returns the final inertia (sum of distances to assigned centroids). - fn lloyd( - &mut self, - sample: &[f32], - chunk: usize, - row_chunk: usize, - sample_inv_norms: &[f32], + seed: u64, config: &Config, ) -> f32 { - let &mut Self { d, k, .. } = self; - let mut previous_inertia = f32::INFINITY; - let mut inertia = f32::INFINITY; - - for _ in 0..config.max_iters.get() { - inertia = sample - .par_chunks(row_chunk) - .zip(sample_inv_norms.par_chunks(chunk)) - .zip(self.labels.par_chunks_mut(chunk)) - .map(|((points, inv_norms), labels)| { - let mut inertia = 0.0; - let count = labels.len(); - - // SAFETY: each parallel chunk pairs `count` labels with - // `count * d` floats of point data and `count` inv_norms. - // `d` is a multiple of 8 (guaranteed by Dimension). - unsafe { - core::hint::assert_unchecked(points.len() == count * d); - core::hint::assert_unchecked(inv_norms.len() == count); - core::hint::assert_unchecked(d.is_multiple_of(8)); - } - - let mut i = 0; - while i + 4 <= count { - let p0 = &points[i * d..i * d + d]; - let p1 = &points[(i + 1) * d..(i + 1) * d + d]; - let p2 = &points[(i + 2) * d..(i + 2) * d + d]; - let p3 = &points[(i + 3) * d..(i + 3) * d + d]; - - // SAFETY: each point length d, centroids length k*d, - // k > 0, d a multiple of 8 (guaranteed by Dimension). - let nearest = - unsafe { kernel::nearest4(p0, p1, p2, p3, &self.centroids, k, d) }; - - let inv = [ - inv_norms[i], - inv_norms[i + 1], - inv_norms[i + 2], - inv_norms[i + 3], - ]; - for m in 0..4 { - labels[i + m] = nearest[m].0; - inertia += squared_chord_distance(nearest[m].1, inv[m]); - } - i += 4; - } - - while i < count { - let point = &points[i * d..i * d + d]; - // SAFETY: point length d, centroids length k*d, k > 0, - // d mult of 8. - let (label, distance) = - unsafe { nearest_centroid(point, inv_norms[i], &self.centroids, k, d) }; - labels[i] = label; - inertia += distance; - i += 1; - } - - inertia - }) - .sum(); - - self.reset_sums(); - self.reset_counts(); - - for ((point, label), inv_norm) in sample - .chunks_exact(d) - .zip(self.labels.iter().copied()) - .zip(sample_inv_norms.iter().copied()) - { - let cluster = usize::from(label); - let start = cluster * d; - - self.counts[cluster] += 1; - - if inv_norm == 0.0 { - continue; - } - - // SAFETY: `sums[start..start + d]` and `point` both have - // length `d`, and `d` is a multiple of 8 (guaranteed by Dimension). - unsafe { - kernel::add_scaled_into(&mut self.sums[start..start + d], point, inv_norm); - } - } - - for cluster in 0..k { - if self.counts[cluster] == 0 { - continue; - } - - let start = cluster * d; - let centroid = &mut self.centroids[start..start + d]; - let sum = &self.sums[start..start + d]; - - // SAFETY: `centroid` and `sum` both have length `D`, and `D` - // is a multiple of 8 (guaranteed by Dimension). - unsafe { - #[expect( - clippy::cast_precision_loss, - reason = "cluster count is bounded by sample_cap (≤8192), well within f32 \ - precision" - )] - let inv_count = 1.0 / self.counts[cluster] as f32; - kernel::scale_into(centroid, sum, inv_count); - } - - // SAFETY: centroid rows have length `D`, and `D` is a - // multiple of 8 (guaranteed by Dimension). - unsafe { - kernel::normalize(centroid); - } - } - - let reseeded = self.reinit_empty_clusters(sample, sample_inv_norms); - - // Skip the convergence check when a cluster was just reseeded: - // the reseeded centroid hasn't had an assignment pass yet, so - // breaking now would waste the reinit. - if !reseeded && previous_inertia.is_finite() { - let relative_change = - (previous_inertia - inertia).abs() / previous_inertia.max(f32::EPSILON); - - if relative_change <= config.tol { - break; - } - } + let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); - previous_inertia = inertia; - } + // `new` zeroes everything else; the distance scratch must start at + // infinity so the first seeding pass overwrites every entry. + self.point_distances.fill(f32::INFINITY); - inertia + self.seed_plusplus(sample, sample_inv_norms, &mut rng); + self.lloyd(sample, chunk, row_chunk, sample_inv_norms, config) } /// k-means++ D² weighted seeding. Picks `k` initial centroids from the /// sample, each chosen with probability proportional to its squared /// distance from the nearest already-chosen centroid. - fn seed_plusplus(&mut self, sample: &[f32], sample_inv_norms: &[f32], mut rng: impl Rng) { + fn seed_plusplus(&mut self, sample: &[f32], sample_inv_norms: &[f32], rng: &mut impl Rng) { let &mut Self { d, k, m, .. } = self; + let mut point = rng.random_range(0..m); - let mut restart_rng = Xoshiro256PlusPlus::seed_from_u64(rng.random()); - let mut point = restart_rng.random_range(0..m); - - for cluster in 0..k { + for cluster in 0..k.get() { let centroid_start = cluster * d; let point_start = point * d; self.centroids[centroid_start..centroid_start + d] .copy_from_slice(&sample[point_start..point_start + d]); - // SAFETY: centroid rows have length `D`, and `D` is a multiple of 8 (guaranteed by - // Dimension). + // SAFETY: centroid rows have length `d`, and `d` is a multiple of 8. unsafe { kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); } self.selected[point] = true; + // The last centroid needs no D² update: those distances would + // only be used to sample a further seed. + if cluster + 1 == k.get() { + break; + } + let centroid = &self.centroids[centroid_start..centroid_start + d]; - let total: f32 = sample + // Per-element writes only, so the pass is deterministic under + // rayon; the D² total is summed sequentially below. + sample .par_chunks_exact(d) - .zip(sample_inv_norms.par_iter().copied()) - .zip(self.closest_distances.par_iter_mut()) + .zip(sample_inv_norms.par_iter()) + .zip(self.point_distances.par_iter_mut()) .enumerate() - .map(|(index, ((point, inv_norm), closest))| { + .for_each(|(index, ((point, &inv_norm), closest))| { if self.selected[index] { *closest = 0.0; - return 0.0; + return; } // SAFETY: `point` and `centroid` both have length `D`, and @@ -552,40 +429,37 @@ impl Fit { if distance < *closest { *closest = distance; } + }); - *closest - }) - .sum(); - - if cluster + 1 == k { - break; - } + let total: f32 = self.point_distances.iter().sum(); point = if total.is_finite() && total > 0.0 { - let mut target = restart_rng.random_range(0.0..total); - let mut sampled = self - .closest_distances - .iter() - .rposition(|distance| *distance > 0.0) - .unwrap_or(0); - - for (index, distance) in self.closest_distances.iter().copied().enumerate() { + let mut target = rng.random_range(0.0..total); + let mut sampled = None; + let mut last_positive = 0; + + for (index, &distance) in self.point_distances.iter().enumerate() { if distance <= 0.0 { continue; } + last_positive = index; target -= distance; if target <= 0.0 { - sampled = index; + sampled = Some(index); break; } } - sampled + // Rounding can leave `target` marginally positive after the last bucket; + // fall back to the last point with positive mass. + sampled.unwrap_or(last_positive) } else { + // Degenerate geometry: every remaining point coincides with a seed. + // Pick uniformly among the unselected points. let remaining = self.selected.iter().filter(|selected| !**selected).count(); - let mut target = restart_rng.random_range(0..remaining); + let mut target = rng.random_range(0..remaining); let mut sampled = 0; for (index, selected) in self.selected.iter().copied().enumerate() { @@ -605,156 +479,234 @@ impl Fit { }; } } -} -/// Per-thread accumulator for parallel centroid recomputation. -/// -/// Each rayon task gets its own `Accum`; they are merged via [`Accum::merge`] -/// after the parallel fold completes. -struct Accum { - /// Per-cluster sum of normalized points, `k * d` elements. - sums: Box<[f32]>, - /// Per-cluster point count. - counts: Box<[usize]>, -} + /// Runs Lloyd iterations on the sample until convergence or `max_iters`. + /// + /// Returns the final inertia (sum of distances to assigned centroids). + fn lloyd( + &mut self, + sample: &[f32], + chunk: usize, + row_chunk: usize, + sample_inv_norms: &[f32], + config: &Config, + ) -> f32 { + let &mut Self { d, k, .. } = self; + let mut previous_inertia = f32::INFINITY; + let mut inertia = f32::INFINITY; -impl Accum { - fn new(k: usize, d: usize) -> Self { - // SAFETY: all-zero bits are valid for f32 (0.0) and usize (0). `assume_init` is - // sound after `new_zeroed_slice`. - let sums = unsafe { Box::<[f32]>::new_zeroed_slice(k * d).assume_init() }; - // SAFETY: see above - let counts = unsafe { Box::<[usize]>::new_zeroed_slice(k).assume_init() }; - Self { sums, counts } + for _ in 0..config.max_iters.get() { + // Assignment: labels and per-point distances. + // Trivially deterministic under rayon, as writes are per element. + sample + .par_chunks(row_chunk) + .zip(sample_inv_norms.par_chunks(chunk)) + .zip(self.labels.par_chunks_mut(chunk)) + .zip(self.point_distances.par_chunks_mut(chunk)) + .for_each(|(((points, inv_norms), labels), distances)| { + // SAFETY: `par_chunks(row_chunk)` with `row_chunk == chunk * d` + // pairs `labels.len()` labels, distances, and inv norms with + // `labels.len() * d` floats of points. `self.centroids` has + // length `k * d`, and `d` is a multiple of 8 (guaranteed by + // Dimension). + unsafe { + lloyd_assign(k, d, &self.centroids, points, inv_norms, labels, distances); + }; + }); + + inertia = self.point_distances.iter().sum(); + + // SAFETY: `sample.len() == m * d` with `m` labels, sums is + // `k * d` with `k` counts, and `d` is a multiple of 8 + // (guaranteed by Dimension). + unsafe { + accumulate_clusters( + sample, + &self.labels, + Some(sample_inv_norms), + &mut self.sums, + &mut self.counts, + d, + ); + } + + for cluster in 0..k.get() { + if self.counts[cluster] == 0 { + continue; + } + + let start = cluster * d; + + // Normalization is scale-invariant, so the raw sum gives the same direction as the + // average. + self.centroids[start..start + d].copy_from_slice(&self.sums[start..start + d]); + + // SAFETY: centroid rows have length `d`, and `d` is a multiple of 8 + // (guaranteed by Dimension). + unsafe { + kernel::normalize(&mut self.centroids[start..start + d]); + } + } + + let reseeded = self.reinit_empty_clusters(sample); + + // Skip the convergence check when a cluster was just reseeded: + // the reseeded centroid hasn't had an assignment pass yet, so + // breaking now would waste the reinit. + if !reseeded && previous_inertia.is_finite() { + let relative_change = + (previous_inertia - inertia).abs() / previous_inertia.max(f32::EPSILON); + + if relative_change <= config.tol { + break; + } + } + + previous_inertia = inertia; + } + + inertia } - fn merge(mut self, other: &Self, k: usize, d: usize) -> Self { - for cluster in 0..k { - let start = cluster * d; + /// Reinitializes empty clusters from the sample point farthest from its + /// assigned centroid, using the distances stored by the assignment pass. + /// + /// After relocating a point its stored distance is zeroed and its label + /// updated, so subsequent empty clusters in the same pass pick different + /// points. + #[expect( + clippy::cast_possible_truncation, + reason = "cluster index < k, and k originates from Config::k (u16)" + )] + fn reinit_empty_clusters(&mut self, sample: &[f32]) -> bool { + let &mut Self { d, k, .. } = self; + let mut reseeded = false; + + for cluster in 0..k.get() { + if self.counts[cluster] != 0 { + continue; + } + + reseeded = true; - self.counts[cluster] += other.counts[cluster]; + let mut farthest_idx = 0; + let mut farthest_dist = -1.0_f32; - // SAFETY: both cluster sum rows have length `d`, and `d` is a - // multiple of 8 (guaranteed by Dimension). + for (index, &distance) in self.point_distances.iter().enumerate() { + if distance > farthest_dist { + farthest_dist = distance; + farthest_idx = index; + } + } + + let point_start = farthest_idx * d; + let centroid_start = cluster * d; + + self.centroids[centroid_start..centroid_start + d] + .copy_from_slice(&sample[point_start..point_start + d]); + + // SAFETY: centroid rows have length `d`, a multiple of 8. unsafe { - kernel::add_into( - &mut self.sums[start..start + d], - &other.sums[start..start + d], - ); + kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); } + + self.labels[farthest_idx] = cluster as u16; + self.point_distances[farthest_idx] = 0.0; } - self + reseeded } } -/// Assigns all `n` points to their nearest centroid, recomputes centroids -/// from the full population, and re-assigns labels to the final centroids. +/// Recomputes per-cluster sums and counts from labeled points. +/// +/// The result is impervious to any thread schedule order. /// -/// Uses a parallel fold/reduce: each rayon task accumulates into its own -/// [`Accum`], then results are merged. The final centroids are averaged -/// and normalized in-place. +/// `inv_norms` supplies precomputed inverse norms; pass `None` to compute +/// them on the fly. +/// +/// Zero-norm points are counted but contribute nothing to the sums. /// /// # Safety /// -/// * `x.len() == n * D` for some `n` -/// * `clustering.centroids.len() == k * D` -/// * `clustering.labels.len() == n` -/// * `k > 0` -/// * `D` is a multiple of 8 (guaranteed by Dimension) -unsafe fn assign(x: &[f32], clustering: &mut Clustering, k: usize, chunk: usize, row_chunk: usize) { - let d = clustering.dimension.get() as usize; - - let full = x - .par_chunks(row_chunk) - .zip(clustering.labels.par_chunks_mut(chunk)) - .fold( - || Accum::new(k, d), - |mut accum, (points, labels)| { - // SAFETY: `cluster` established `x.len() == n * D` and - // `centroids.len() == k * D`. `par_chunks(row_chunk)` with - // `row_chunk = chunk * D` produces chunks where - // `points.len()` is a multiple of `D` and matches `labels.len() * D`. - unsafe { - assign_chunk(&clustering.centroids, k, d, points, labels, &mut accum); +/// * `points.len() == labels.len() * d` +/// * `sums.len() == counts.len() * d` +/// * `inv_norms`, when provided, has one entry per label +/// * `d` is a multiple of 8 +unsafe fn accumulate_clusters( + points: &[f32], + labels: &[u16], + inv_norms: Option<&[f32]>, + sums: &mut [f32], + counts: &mut [usize], + d: usize, +) { + // A `debug_assert` only: an `assert_unchecked` here would be sound (the + // length is a documented precondition), but to elide the + // `inv_norms[index]` bounds check the fact would have to survive into + // the rayon closure, and that check is noise next to the `d`-wide + // kernel call it precedes. + debug_assert!(inv_norms.is_none_or(|norms| norms.len() == labels.len())); + + sums.par_chunks_exact_mut(d) + .zip(counts.par_iter_mut()) + .enumerate() + .for_each(|(cluster, (sum, count))| { + sum.fill(0.0); + *count = 0; + + for (index, (point, &label)) in points.chunks_exact(d).zip(labels).enumerate() { + if usize::from(label) != cluster { + continue; } - accum - }, - ) - .reduce(|| Accum::new(k, d), |lhs, rhs| lhs.merge(&rhs, k, d)); - - for cluster in 0..k { - if full.counts[cluster] == 0 { - continue; - } + *count += 1; - let start = cluster * d; + let inv_norm = inv_norms.map_or_else( + || { + // SAFETY: `point` has length `d`, a multiple of 8 + // (guaranteed by the caller). + let norm = unsafe { kernel::dot(point, point) }.sqrt(); - #[expect( - clippy::cast_possible_truncation, - reason = "cluster < k and k originates from Config::k (u16)" - )] - let centroid = clustering.centroid_mut(cluster as u16); - let sum = &full.sums[start..start + d]; + if norm > 0.0 { norm.recip() } else { 0.0 } + }, + |inv_norms| inv_norms[index], + ); - // SAFETY: centroid and sum both length D, a multiple of 8. - unsafe { - #[expect( - clippy::cast_precision_loss, - reason = "cluster count bounded by n; precision loss acceptable for averaging" - )] - let inv_count = 1.0 / full.counts[cluster] as f32; - kernel::scale_into(centroid, sum, inv_count); - } - // SAFETY: centroid length D, a multiple of 8. - unsafe { - kernel::normalize(centroid); - } - } + if inv_norm == 0.0 { + continue; + } - // SAFETY: centroids were just recomputed; same invariants hold. - unsafe { - reassign( - x, - &clustering.centroids, - &mut clustering.labels, - k, - d, - chunk, - row_chunk, - ); - } + // SAFETY: `sum` and `point` both have length `d`, and `d` is + // a multiple of 8 (guaranteed by the caller). + unsafe { + kernel::add_scaled_into(sum, point, inv_norm); + } + } + }); } -/// Processes one parallel chunk of the assignment step: finds the nearest -/// centroid for each point, accumulates normalized points into cluster sums, -/// and records labels. +/// Labels one parallel chunk: each point gets its nearest centroid. /// /// # Safety /// /// * `points.len() == labels.len() * d` /// * `centroids.len() >= k * d` /// * `d` is a multiple of 8 -/// * `k > 0` -/// * `accum.sums.len() >= k * d` and `accum.counts.len() >= k` -unsafe fn assign_chunk( +unsafe fn label_chunk( centroids: &[f32], - k: usize, + k: NonZero, d: usize, points: &[f32], labels: &mut [u16], - accum: &mut Accum, ) { - // field path -> disjoint capture of `centroids` only, leaving - // `labels` free for the mutable parallel borrow. let count = labels.len(); - // SAFETY: each parallel chunk pairs `count` labels with - // `count * D` floats of point data. `D` is a compile-time - // multiple of 8. + // SAFETY: each parallel chunk pairs `count` labels with `count * d` + // floats of point data; `d` is a multiple of 8 (guaranteed by Dimension). unsafe { core::hint::assert_unchecked(points.len() == count * d); + core::hint::assert_unchecked(centroids.len() >= k.get() * d); core::hint::assert_unchecked(d.is_multiple_of(8)); } @@ -765,83 +717,51 @@ unsafe fn assign_chunk( let p2 = &points[(i + 2) * d..(i + 2) * d + d]; let p3 = &points[(i + 3) * d..(i + 3) * d + d]; - // SAFETY: each point length D, centroids length k*D, k > 0, D a multiple of 8 (guaranteed - // by Dimension). + // SAFETY: each point length d, centroids length k*d, k > 0, + // d a multiple of 8 (guaranteed by Dimension). let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; - let ps = [p0, p1, p2, p3]; - - for m in 0..4 { - let label = nearest[m].0; - labels[i + m] = label; - let cluster = usize::from(label); - accum.counts[cluster] += 1; - let start = cluster * d; - - // SAFETY: point length D, a multiple of 8. - let norm = unsafe { kernel::dot(ps[m], ps[m]).sqrt() }; - if norm == 0.0 { - continue; - } - - // SAFETY: `sums[start..start + D]` and `point` both have length `D`, and `D` is a - // multiple of 8 (guaranteed by Dimension). - unsafe { - kernel::add_scaled_into(&mut accum.sums[start..start + d], ps[m], norm.recip()); - } - } + labels[i] = nearest[0].0; + labels[i + 1] = nearest[1].0; + labels[i + 2] = nearest[2].0; + labels[i + 3] = nearest[3].0; i += 4; } while i < count { let point = &points[i * d..i * d + d]; - // SAFETY: point length D, centroids length k*D, k > 0, D mult of 8. + // SAFETY: point length d, centroids length k*d, k > 0, d mult of 8. let (label, _) = unsafe { nearest_centroid(point, 1.0, centroids, k, d) }; labels[i] = label; - let cluster = usize::from(label); - accum.counts[cluster] += 1; - - let start = cluster * d; - - // SAFETY: point length D. - let norm = unsafe { kernel::dot(point, point).sqrt() }; - if norm != 0.0 { - // SAFETY: `sums[start..start + D]` and `point` both - // have length `D`, and `D` is a multiple of 8. - unsafe { - kernel::add_scaled_into(&mut accum.sums[start..start + d], point, norm.recip()); - } - } i += 1; } } -/// Processes one parallel chunk of the reassignment step: updates each -/// label to the nearest final centroid. +/// Labels one parallel chunk against the final centroids and returns its +/// inertia contribution. Inverse norms are computed on the fly. /// /// # Safety /// /// * `points.len() == labels.len() * d` /// * `centroids.len() >= k * d` /// * `d` is a multiple of 8 -/// * `k > 0` -unsafe fn reassign_chunk( - k: usize, - d: usize, +unsafe fn score_chunk( centroids: &[f32], + k: NonZero, + d: usize, points: &[f32], labels: &mut [u16], -) { - let count = labels.len(); +) -> f32 { + debug_assert_eq!(points.len(), labels.len() * d); - // SAFETY: each parallel chunk pairs `count` labels with - // `count * D` floats of point data. `D` is a compile-time - // multiple of 8. + // SAFETY: The caller must ensure `points.len() == labels.len() * d`. unsafe { - core::hint::assert_unchecked(points.len() == count * d); - core::hint::assert_unchecked(d.is_multiple_of(8)); + core::hint::assert_unchecked(points.len() == labels.len() * d); } + let count = labels.len(); + let mut inertia = 0.0_f32; + let mut i = 0; while i + 4 <= count { let p0 = &points[i * d..i * d + d]; @@ -849,39 +769,52 @@ unsafe fn reassign_chunk( let p2 = &points[(i + 2) * d..(i + 2) * d + d]; let p3 = &points[(i + 3) * d..(i + 3) * d + d]; - // SAFETY: each point length D, centroids length k*D, k > 0, - // D a multiple of 8 (guaranteed by Dimension). + // SAFETY: each point length d, centroids length k*d, k > 0, + // d a multiple of 8 (guaranteed by Dimension). let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; + let ps = [p0, p1, p2, p3]; - labels[i] = nearest[0].0; - labels[i + 1] = nearest[1].0; - labels[i + 2] = nearest[2].0; - labels[i + 3] = nearest[3].0; + for offset in 0..4 { + labels[i + offset] = nearest[offset].0; + + // SAFETY: point length d, a multiple of 8. + let norm = unsafe { kernel::dot(ps[offset], ps[offset]) }.sqrt(); + let inv_norm = if norm > 0.0 { norm.recip() } else { 0.0 }; + inertia += squared_chord_distance(nearest[offset].1, inv_norm); + } i += 4; } while i < count { let point = &points[i * d..i * d + d]; - // SAFETY: point length D, centroids length k*D, k > 0, D mult of 8. - let (label, _) = unsafe { nearest_centroid(point, 1.0, centroids, k, d) }; + + // SAFETY: point length d, a multiple of 8. + let norm = unsafe { kernel::dot(point, point) }.sqrt(); + let inv_norm = if norm > 0.0 { norm.recip() } else { 0.0 }; + + // SAFETY: point length d, centroids length k*d, k > 0, d mult of 8. + let (label, distance) = unsafe { nearest_centroid(point, inv_norm, centroids, k, d) }; labels[i] = label; + inertia += distance; i += 1; } + + inertia } -/// Re-assigns labels to the nearest final centroid. -/// -/// After centroid recomputation, some boundary points may no longer be -/// nearest to the centroid stored under their label. This pass fixes that. +/// Labels every point with its nearest centroid. /// /// # Safety /// -/// Same as [`assign`]. +/// * `x.len() == n * d` for some `n` +/// * `clustering.centroids.len() == k * d` +/// * `clustering.labels.len() == n` +/// * `d` is a multiple of 8 unsafe fn reassign( x: &[f32], centroids: &[f32], labels: &mut [u16], - k: usize, + k: NonZero, d: usize, chunk: usize, row_chunk: usize, @@ -889,19 +822,149 @@ unsafe fn reassign( x.par_chunks(row_chunk) .zip(labels.par_chunks_mut(chunk)) .for_each(|(points, labels)| { - // SAFETY: `par_chunks(row_chunk)` with `row_chunk = chunk * D` - // ensures `points.len() == labels.len() * D`. Centroids and k + // SAFETY: `par_chunks(row_chunk)` with `row_chunk = chunk * d` + // ensures `points.len() == labels.len() * d`. Centroids and k // are valid from the caller. unsafe { - reassign_chunk(k, d, centroids, points, labels); + label_chunk(centroids, k, d, points, labels); } }); } +/// Labels every point with its nearest centroid and returns the total +/// inertia. +/// +/// Labels are exact; the inertia is precise only up to floating-point +/// summation order. +/// +/// # Safety +/// +/// * `x.len() == n * d` for some `n` +/// * `clustering.centroids.len() == k * d` +/// * `clustering.labels.len() == n` +/// * `d` is a multiple of 8 +unsafe fn reassign_scored( + x: &[f32], + centroids: &[f32], + labels: &mut [u16], + k: NonZero, + d: usize, + chunk: usize, + row_chunk: usize, +) -> f32 { + // Unordered parallel reduction: the grouping follows rayon's scheduling, so the sum is not + // bit-stable. An ordered reduction would need to collect per-chunk partials, costing an + // allocation per call. + x.par_chunks(row_chunk) + .zip(labels.par_chunks_mut(chunk)) + .map(|(points, labels)| { + // SAFETY: `par_chunks(row_chunk)` with `row_chunk = chunk * d` + // ensures `points.len() == labels.len() * d`. Centroids and k + // are valid from the caller. + unsafe { score_chunk(centroids, k, d, points, labels) } + }) + .sum() +} + +/// Assigns all `n` points to their nearest centroid, recomputes centroids +/// from the full population, and re-labels against the final centroids. +/// Returns the full-data inertia. +/// +/// `sums` and `counts` are accumulator scratch; their contents on entry are +/// irrelevant. +/// +/// # Safety +/// +/// * `x.len() == n * d` for some `n` +/// * `clustering.centroids.len() == k * d` +/// * `clustering.labels.len() == n` +/// * `sums.len() == k * d` and `counts.len() == k` +/// * `d` is a multiple of 8 +unsafe fn assign( + x: &[f32], + clustering: &mut Clustering, + k: NonZero, + chunk: usize, + row_chunk: usize, + sums: &mut [f32], + counts: &mut [usize], +) -> f32 { + let d = clustering.dimension.get() as usize; + + // 1. Label all points against the sample-fitted centroids. + // SAFETY: forwarded from the caller. + unsafe { + reassign( + x, + &clustering.centroids, + &mut clustering.labels, + k, + d, + chunk, + row_chunk, + ); + } + + // 2. Recompute centroids from the full population. + // SAFETY: `x.len() == n * d` with `n` labels, sums is `k * d` with `k` + // counts, and `d` is a multiple of 8 (guaranteed by Dimension). + unsafe { + accumulate_clusters(x, &clustering.labels, None, sums, counts, d); + } + + for (cluster, count) in counts.iter_mut().enumerate() { + if *count == 0 { + continue; + } + + let start = cluster * d; + + #[expect( + clippy::cast_possible_truncation, + reason = "cluster < k and k originates from Config::k (u16)" + )] + let centroid = clustering.centroid_mut(cluster as u16); + // Normalization is scale-invariant, so the raw sum gives the same + // direction as the average. + centroid.copy_from_slice(&sums[start..start + d]); + + // SAFETY: centroid length d, a multiple of 8. + unsafe { + kernel::normalize(centroid); + } + } + + // 3. Final labels and inertia against the recomputed centroids. + // SAFETY: centroids were just recomputed in place; same invariants hold. + unsafe { + reassign_scored( + x, + &clustering.centroids, + &mut clustering.labels, + k, + d, + chunk, + row_chunk, + ) + } +} + /// Runs spherical k-means over a flat row-major embedding matrix. /// /// `x` contains `n` points of `dimension` floats each, laid out -/// contiguously. Returns cluster assignments and unit-normalized centroids. +/// contiguously. Returns cluster assignments, unit-normalized centroids, and +/// the full-data inertia. +/// +/// Given the same input and configuration, the returned labels and +/// centroids are identical across runs; the inertia is precise only up to +/// floating-point summation order (see [`Clustering::inertia`]). +/// +/// Zero-norm points do not influence centroids, and are always assigned to +/// cluster 0 at distance 0. If a cluster consists solely of zero-norm points, +/// its centroid is zero; see [`Clustering::centroids`]. +/// +/// If `config.k == 0` or `x` is empty there is nothing to fit: the result +/// has no centroids, all-zero placeholder labels, and an inertia of `0.0`. /// /// # Panics /// @@ -917,21 +980,21 @@ pub fn cluster(x: &[f32], dimension: Dimension, config: &Config) -> Clustering { let mut clustering = Clustering::new(k, n, dimension); - if k == 0 { + let Some(k) = NonZero::new(k) else { return clustering; - } + }; - let k = usize::from(k); + let k = NonZero::from(k); let mut rng = Xoshiro256PlusPlus::seed_from_u64(config.seed); // 1. subsample (fit on all of n only when n is already small) - let m = config.sample_cap.max(k).min(n); + let m = config.sample_cap.max(k.get()).min(n); let sample = if m == n { Cow::Borrowed(x) } else { let indices = sample_indices(n, m, &mut rng); - let mut sampled = vec![0_f32; m * d]; + let mut sampled = vec![0.0_f32; m * d]; let chunks = sampled.chunks_mut(d); assert_eq!(chunks.len(), indices.len()); @@ -944,34 +1007,71 @@ pub fn cluster(x: &[f32], dimension: Dimension, config: &Config) -> Clustering { }; let sample = sample.as_ref(); - let chunk = config.chunk.get(); - let row_chunk = chunk - .checked_mul(d) - .unwrap_or_else(|| usize::MAX - (usize::MAX % d)) - .max(d); + + // Clamping to `n` keeps `row_chunk` from overflowing: `chunk * d` is at most `n * d == + // x.len()`. + let chunk = cmp::min(config.chunk.get(), n); + let row_chunk = chunk * d; let sample_inv_norms: Vec = sample .par_chunks_exact(d) .map(|point| { // SAFETY: every point is a `d`-sized row, and `d` is a multiple of 8 (guaranteed by // Dimension). - let norm = unsafe { kernel::dot(point, point).sqrt() }; + let norm = unsafe { kernel::dot(point, point) }.sqrt(); if norm > 0.0 { norm.recip() } else { 0.0 } }) .collect(); - // 2. fit on the sample, best of n_init restarts (guards against bad initializations) - let mut fit = Fit::new(k, m, d); - fit.run(sample, chunk, row_chunk, &sample_inv_norms, rng, config); - mem::swap(&mut clustering.centroids, &mut fit.best_centroids); + // 2. fit on the sample: independent k-means++ restarts in parallel, the + // run with the lowest inertia wins (guards against bad initializations). + // Seeds are pre-derived so the stream matches a sequential run; ties + // break on the restart index, which keeps the winner deterministic no + // matter how rayon schedules the restarts. + let seeds: Vec = core::iter::repeat_with(|| rng.random()) + .take(usize::try_from(config.n_init.get()).unwrap_or(usize::MAX)) + .collect(); + + let best = seeds + .into_par_iter() + .enumerate() + .map(|(index, seed)| { + let mut restart = Restart::new(k, m, d); + let inertia = restart.run(sample, chunk, row_chunk, &sample_inv_norms, seed, config); + + (inertia, index, restart) + }) + .min_by(|lhs, rhs| lhs.0.total_cmp(&rhs.0).then(lhs.1.cmp(&rhs.1))) + .expect("config.n_init is non-zero, so at least one restart ran"); + + // Reuse the winning restart's buffers: its centroids become the result + // and its per-cluster accumulators serve the full-data recomputation, + // instead of allocating fresh ones. + let Restart { + centroids, + mut sums, + mut counts, + .. + } = best.2; + + clustering.centroids = centroids; // 3. assign points to clusters // SAFETY: `x.len() == n * d` (asserted above), `clustering.centroids.len() == k * d`, - // `k > 0` (checked above), `d` is a multiple of 8 (guaranteed by Dimension). - unsafe { - assign(x, &mut clustering, k, chunk, row_chunk); - } + // `sums` and `counts` are the restart's `k * d` and `k` sized accumulators, + // and `d` is a multiple of 8 (guaranteed by Dimension). + clustering.inertia = unsafe { + assign( + x, + &mut clustering, + k, + chunk, + row_chunk, + &mut sums, + &mut counts, + ) + }; clustering } @@ -986,6 +1086,12 @@ mod tests { )] use super::*; + macro_rules! nz { + ($expr:expr) => { + const { NonZero::new($expr).unwrap() } + }; + } + /// Builds well-separated blob clusters in D-dimensional space. /// /// Each blob has a dominant axis so clusters are far apart in cosine @@ -1029,9 +1135,9 @@ mod tests { } /// Random unit-norm centroids in `D`-dimensional space. - fn unit_random(k: usize, seed: u64) -> Vec { + fn unit_random(k: NonZero, seed: u64) -> Vec { let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); - let mut c = vec![0.0_f32; k * D]; + let mut c = vec![0.0_f32; k.get() * D]; for row in c.chunks_exact_mut(D) { for v in row.iter_mut() { *v = rng.random_range(-1.0..1.0); @@ -1046,11 +1152,12 @@ mod tests { /// Brute-force nearest centroid by cosine similarity. #[expect(clippy::cast_possible_truncation, reason = "k is small in tests")] - fn brute_nearest_cosine(point: &[f32], centroids: &[f32], k: usize) -> u16 { + fn brute_nearest_cosine(point: &[f32], centroids: &[f32], k: NonZero) -> u16 { let pn = l2(point); let mut best = 0_u16; let mut best_cos = f32::NEG_INFINITY; - for c in 0..k { + + for c in 0..k.get() { let cent = ¢roids[c * D..(c + 1) * D]; let d: f32 = point.iter().zip(cent).map(|(a, b)| a * b).sum(); let cn = l2(cent); @@ -1059,6 +1166,7 @@ mod tests { } else { d / (pn * cn) }; + if cos > best_cos { best_cos = cos; best = c as u16; @@ -1131,12 +1239,26 @@ mod tests { } } + #[test] + fn sample_indices_unique_sorted_in_range() { + let rng = Xoshiro256PlusPlus::seed_from_u64(1); + let indices = sample_indices(1000, 100, rng); + + assert_eq!(indices.len(), 100); + assert!( + indices.is_sorted_by(|lhs, rhs| lhs < rhs), + "indices must be strictly increasing (sorted, unique)" + ); + assert!(indices.iter().all(|&index| index < 1000)); + } + #[test] fn cluster_empty_input() { let config = Config::for_k_with_seed(4, 42); let result = cluster(&[], dim(8), &config); assert_eq!(result.labels.len(), 0); assert_eq!(result.centroids.len(), 0); + assert_eq!(result.inertia, 0.0); } #[test] @@ -1146,6 +1268,8 @@ mod tests { let result = cluster(&data, dim(8), &config); assert_eq!(result.labels.len(), 1); assert_eq!(result.labels[0], 0); + assert_eq!(result.centroids.len(), 0); + assert_eq!(result.inertia, 0.0); } #[test] @@ -1229,21 +1353,31 @@ mod tests { assert_eq!(r1.labels, r2.labels); assert_eq!(r1.centroids, r2.centroids); + + // The inertia reduction is a parallel sum, so it is only + // deterministic up to float summation order. + let tolerance = r1.inertia.abs().max(f32::EPSILON) * 1e-5; + assert!( + (r1.inertia - r2.inertia).abs() <= tolerance, + "inertia should agree within summation-order tolerance: {} vs {}", + r1.inertia, + r2.inertia + ); } #[test] - fn cluster_different_seeds_may_differ() { - let (data, _) = make_blobs::<8>(30, 3, 555); + fn cluster_recovers_blobs_across_seeds() { + let (data, truth) = make_blobs::<8>(30, 3, 555); - let r1 = cluster(&data, dim(8), &Config::for_k_with_seed(3, 42)); - let r2 = cluster(&data, dim(8), &Config::for_k_with_seed(3, 9999)); - - // Not guaranteed to differ, but with well-separated blobs and - // different seeds the label permutation usually differs. - assert!( - r1.labels != r2.labels, - "different seeds produced identical label vectors (possible but unlikely)" - ); + for seed in [42, 9999] { + let result = cluster(&data, dim(8), &Config::for_k_with_seed(3, seed)); + let acc = accuracy(&result.labels, &truth, 3); + assert!( + acc > 0.95, + "seed {seed}: expected >95% accuracy, got {:.1}%", + acc * 100.0 + ); + } } #[test] @@ -1314,6 +1448,72 @@ mod tests { ); } + #[test] + fn cluster_d256_recovers_blobs() { + // Production default dimension (matryoshka truncation target). + let (data, truth) = make_blobs::<256>(50, 4, 1234); + let config = Config::for_k_with_seed(4, 42); + let result = cluster(&data, dim(256), &config); + + let acc = accuracy(&result.labels, &truth, 4); + assert!( + acc > 0.95, + "D=256: expected >95% accuracy, got {:.1}%", + acc * 100.0 + ); + } + + #[test] + fn cluster_d1536_recovers_blobs() { + let (data, truth) = make_blobs::<1536>(20, 3, 4321); + let config = Config::for_k_with_seed(3, 42); + let result = cluster(&data, dim(1536), &config); + + let acc = accuracy(&result.labels, &truth, 3); + assert!( + acc > 0.95, + "D=1536: expected >95% accuracy, got {:.1}%", + acc * 100.0 + ); + } + + #[test] + fn cluster_chunk_sizes_produce_valid_results() { + let (data, truth) = make_blobs::<8>(50, 4, 314); + + for chunk in [1_usize, 3, 1_000_000] { + let mut config = Config::for_k_with_seed(4, 42); + config.chunk = NonZero::new(chunk).expect("chunk is non-zero"); + let result = cluster(&data, dim(8), &config); + + let acc = accuracy(&result.labels, &truth, 4); + assert!( + acc > 0.95, + "chunk={chunk}: expected >95% accuracy, got {:.1}%", + acc * 100.0 + ); + } + } + + #[test] + fn cluster_inertia_reflects_fit_quality() { + let (data, _) = make_blobs::<8>(50, 4, 99); + + let tight = cluster(&data, dim(8), &Config::for_k_with_seed(4, 42)); + assert!(tight.inertia.is_finite()); + assert!(tight.inertia >= 0.0); + + // Forcing 4 well-separated blobs into a single cluster must fit + // strictly worse. + let loose = cluster(&data, dim(8), &Config::for_k_with_seed(1, 42)); + assert!( + loose.inertia > tight.inertia, + "k=1 inertia {} should exceed k=4 inertia {}", + loose.inertia, + tight.inertia + ); + } + #[test] fn cluster_recovers_with_subsampling() { // n=12000 with sample_cap=1024 exercises the Cow::Owned path. @@ -1361,7 +1561,7 @@ mod tests { #[test] fn nearest_centroid_matches_brute_force_cosine() { - let k = 7; + let k = nz!(7); let centroids = unit_random(k, 99); let mut rng = Xoshiro256PlusPlus::seed_from_u64(100); @@ -1385,7 +1585,7 @@ mod tests { #[test] fn nearest_centroid_argmax_independent_of_inv_norm() { - let k = 5; + let k = nz!(5); let centroids = unit_random(k, 7); let mut rng = Xoshiro256PlusPlus::seed_from_u64(8); diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/store/src/embedding/kernel.rs index 1b971bc23b1..67f99a66831 100644 --- a/libs/@local/graph/store/src/embedding/kernel.rs +++ b/libs/@local/graph/store/src/embedding/kernel.rs @@ -3,7 +3,10 @@ reason = "while usually discouraged, SIMD operations need to be inlined, as otherwise we \ spill SIMD registers, see the SIMD documentation." )] -use core::simd::{Simd, f32x8, num::SimdFloat as _}; +use core::{ + num::NonZero, + simd::{Simd, f32x8, num::SimdFloat as _}, +}; /// Fused multiply-add when the target has native FMA, separate mul+add otherwise. /// @@ -39,7 +42,7 @@ fn simd_mul_add(lhs: f32x8, rhs: f32x8, acc: f32x8) -> f32x8 { /// * Both lengths are multiples of 8. #[inline] #[must_use] -pub(crate) unsafe fn dot(lhs: &[f32], rhs: &[f32]) -> f32 { +pub unsafe fn dot(lhs: &[f32], rhs: &[f32]) -> f32 { debug_assert!(lhs.len().is_multiple_of(8) && lhs.len() == rhs.len()); // SAFETY: the caller guarantees equal lengths and a multiple of 8. @@ -96,63 +99,6 @@ pub(crate) unsafe fn dot(lhs: &[f32], rhs: &[f32]) -> f32 { (s0 + s1 + s2 + s3).reduce_sum() } -/// Adds `src` element-wise into `dst`. -/// -/// # Safety -/// -/// * `dst.len() == src.len()` -/// * Both lengths are multiples of 8. -#[inline] -pub(crate) unsafe fn add_into(dst: &mut [f32], src: &[f32]) { - debug_assert!(dst.len().is_multiple_of(8) && dst.len() == src.len()); - - // SAFETY: the caller guarantees equal lengths and a multiple of 8. - unsafe { - core::hint::assert_unchecked(dst.len() == src.len()); - core::hint::assert_unchecked(dst.len().is_multiple_of(8)); - } - - let (dst, _) = dst.as_chunks_mut::<8>(); - let (src, _) = src.as_chunks::<8>(); - - // SAFETY: same reasoning as the pre-chunk hints: equal input lengths - // that are multiples of 8 produce equal chunk counts. - unsafe { core::hint::assert_unchecked(dst.len() == src.len()) } - - for index in 0..dst.len() { - dst[index] = (f32x8::from_slice(&dst[index]) + f32x8::from_slice(&src[index])).to_array(); - } -} - -/// Writes `src * factor` element-wise into `dst`. -/// -/// # Safety -/// -/// * `dst.len() == src.len()` -/// * Both lengths are multiples of 8. -#[inline] -pub(crate) unsafe fn scale_into(dst: &mut [f32], src: &[f32], factor: f32) { - debug_assert!(dst.len().is_multiple_of(8) && dst.len() == src.len()); - - // SAFETY: the caller guarantees equal lengths and a multiple of 8. - unsafe { - core::hint::assert_unchecked(dst.len() == src.len()); - core::hint::assert_unchecked(dst.len().is_multiple_of(8)); - } - - let factor = f32x8::splat(factor); - let (dst, _) = dst.as_chunks_mut::<8>(); - let (src, _) = src.as_chunks::<8>(); - - // SAFETY: same reasoning as the pre-chunk hints: equal input lengths - // that are multiples of 8 produce equal chunk counts. - unsafe { core::hint::assert_unchecked(dst.len() == src.len()) } - - for index in 0..dst.len() { - dst[index] = (f32x8::from_slice(&src[index]) * factor).to_array(); - } -} - /// Scales `value` in-place by `factor`. /// /// # Safety @@ -185,7 +131,7 @@ pub(crate) unsafe fn scale(value: &mut [f32], factor: f32) { /// * `dst.len() == src.len()` /// * Both lengths are multiples of 8. #[inline] -pub(crate) unsafe fn add_scaled_into(dst: &mut [f32], src: &[f32], factor: f32) { +pub unsafe fn add_scaled_into(dst: &mut [f32], src: &[f32], factor: f32) { debug_assert!(dst.len().is_multiple_of(8) && dst.len() == src.len()); // SAFETY: the caller guarantees equal lengths and a multiple of 8. @@ -247,7 +193,8 @@ pub(crate) unsafe fn normalize(value: &mut [f32]) { /// * all six slices have length `d` /// * `d` is a multiple of 8 #[inline(always)] // micro-kernel must inline nearest4 to keep accumulators in registers -pub(crate) unsafe fn micro_4x2( +#[must_use] +pub unsafe fn micro_4x2( p0: &[f32], p1: &[f32], p2: &[f32], @@ -316,6 +263,67 @@ pub(crate) unsafe fn micro_4x2( ] } +/// 4 points x 1 centroid: the odd-`k` remainder of [`nearest4`]. Four +/// independent accumulators (one per point) share each centroid load, so the +/// centroid streams through registers once instead of once per point. +/// +/// # Safety +/// +/// * all five slices have length `d` +/// * `d` is a multiple of 8 +#[inline(always)] // micro-kernel must inline into nearest4 to keep accumulators in registers +pub(crate) unsafe fn micro_4x1( + p0: &[f32], + p1: &[f32], + p2: &[f32], + p3: &[f32], + c: &[f32], +) -> [f32; 4] { + debug_assert!(p0.len().is_multiple_of(8)); + debug_assert!( + [p1.len(), p2.len(), p3.len(), c.len()] + .iter() + .all(|&l| l == p0.len()) + ); + + let (p0, _) = p0.as_chunks::<8>(); + let (p1, _) = p1.as_chunks::<8>(); + let (p2, _) = p2.as_chunks::<8>(); + let (p3, _) = p3.as_chunks::<8>(); + let (c, _) = c.as_chunks::<8>(); + + // SAFETY: the caller guarantees all five slices have equal length `d`, + // and `d` is a multiple of 8. The hints let the compiler prove that + // `as_chunks` produces equal-length chunk slices. + unsafe { + core::hint::assert_unchecked(p0.len() == p1.len()); + core::hint::assert_unchecked(p0.len() == p2.len()); + core::hint::assert_unchecked(p0.len() == p3.len()); + core::hint::assert_unchecked(p0.len() == c.len()); + } + + let mut a0 = f32x8::splat(0.0); + let mut a1 = f32x8::splat(0.0); + let mut a2 = f32x8::splat(0.0); + let mut a3 = f32x8::splat(0.0); + + for t in 0..c.len() { + let v = Simd::from_array(c[t]); + + a0 = simd_mul_add(Simd::from_array(p0[t]), v, a0); + a1 = simd_mul_add(Simd::from_array(p1[t]), v, a1); + a2 = simd_mul_add(Simd::from_array(p2[t]), v, a2); + a3 = simd_mul_add(Simd::from_array(p3[t]), v, a3); + } + + [ + a0.reduce_sum(), + a1.reduce_sum(), + a2.reduce_sum(), + a3.reduce_sum(), + ] +} + /// Finds the nearest centroid for 4 points simultaneously using the /// [`micro_4x2`] tiled kernel. /// @@ -328,16 +336,15 @@ pub(crate) unsafe fn micro_4x2( /// * All four point slices have length `d`. /// * `centroids.len() >= k * d`. /// * `d` is a multiple of 8. -/// * `k > 0`. #[inline] #[must_use] -pub(crate) unsafe fn nearest4( +pub unsafe fn nearest4( p0: &[f32], p1: &[f32], p2: &[f32], p3: &[f32], centroids: &[f32], - k: usize, + k: NonZero, d: usize, ) -> [(u16, f32); 4] { let mut best_dot = [f32::NEG_INFINITY; 4]; @@ -349,13 +356,12 @@ pub(crate) unsafe fn nearest4( core::hint::assert_unchecked(p0.len() == p1.len()); core::hint::assert_unchecked(p0.len() == p2.len()); core::hint::assert_unchecked(p0.len() == p3.len()); - core::hint::assert_unchecked(centroids.len() >= k * d); + core::hint::assert_unchecked(centroids.len() >= k.get() * d); core::hint::assert_unchecked(d.is_multiple_of(8)); - core::hint::assert_unchecked(k > 0); } let mut j = 0; - while j + 2 <= k { + while j + 2 <= k.get() { // SAFETY: `j + 2 <= k` and `centroids.len() >= k * d`, so both // slices `[j*d .. (j+2)*d]` are in-bounds. let c0 = unsafe { centroids.get_unchecked(j * d..j * d + d) }; @@ -382,19 +388,19 @@ pub(crate) unsafe fn nearest4( j += 2; } - // Handle odd k: one remaining centroid. - if j < k { + // Handle odd k: one remaining centroid via the 4x1 tile. + if j < k.get() { let c = ¢roids[j * d..j * d + d]; - let ps = [p0, p1, p2, p3]; + // SAFETY: all five slices have length `d`, a multiple of 8. + let dots = unsafe { micro_4x1(p0, p1, p2, p3, c) }; + + #[expect( + clippy::cast_possible_truncation, + reason = "k originates from Config::k (u16)" + )] for m in 0..4 { - // SAFETY: point and centroid both have length `d`, a multiple of 8. - let d = unsafe { dot(ps[m], c) }; - #[expect( - clippy::cast_possible_truncation, - reason = "k originates from Config::k (u16)" - )] - if d > best_dot[m] { - best_dot[m] = d; + if dots[m] > best_dot[m] { + best_dot[m] = dots[m]; best_idx[m] = j as u16; } } @@ -414,6 +420,12 @@ mod tests { use super::*; + macro_rules! nz { + ($expr:expr) => { + const { NonZero::new($expr).unwrap() } + }; + } + /// Scalar dot product for reference. fn ref_dot(a: &[f32], b: &[f32]) -> f32 { a.iter().zip(b).map(|(x, y)| x * y).sum() @@ -510,55 +522,6 @@ mod tests { assert_eq!(got, 0.0); } - #[test] - fn add_into_matches_scalar() { - let src = ramp(16, 1.0); - let mut dst = ramp(16, 0.5); - let expected: Vec = dst.iter().zip(&src).map(|(d, s)| d + s).collect(); - // SAFETY: both slices have length 16, a multiple of 8. - unsafe { add_into(&mut dst, &src) } - assert_eq!(dst, expected); - } - - #[test] - fn add_into_zero_is_identity() { - let zeros = vec![0.0_f32; 24]; - let mut dst = ramp(24, 1.0); - let original = dst.clone(); - // SAFETY: both slices have length 24, a multiple of 8. - unsafe { add_into(&mut dst, &zeros) } - assert_eq!(dst, original); - } - - #[test] - fn scale_into_matches_scalar() { - let src = ramp(16, 1.0); - let mut dst = vec![0.0_f32; 16]; - let factor = 2.5; - let expected: Vec = src.iter().map(|x| x * factor).collect(); - // SAFETY: both slices have length 16, a multiple of 8. - unsafe { scale_into(&mut dst, &src, factor) } - assert_eq!(dst, expected); - } - - #[test] - fn scale_into_zero_gives_zeros() { - let src = ramp(8, 1.0); - let mut dst = ramp(8, 999.0); - // SAFETY: both slices have length 8, a multiple of 8. - unsafe { scale_into(&mut dst, &src, 0.0) } - assert!(dst.iter().all(|&x| x == 0.0)); - } - - #[test] - fn scale_into_one_is_copy() { - let src = ramp(16, 0.3); - let mut dst = vec![0.0_f32; 16]; - // SAFETY: both slices have length 16, a multiple of 8. - unsafe { scale_into(&mut dst, &src, 1.0) } - assert_eq!(dst, src); - } - #[test] fn scale_matches_scalar() { let mut v = ramp(16, 1.0); @@ -686,14 +649,62 @@ mod tests { } } + #[test] + fn micro_4x1_matches_individual_dots() { + let d = 16; + let p0 = ramp(d, 0.1); + let p1 = ramp(d, -0.2); + let p2 = ramp(d, 0.3); + let p3 = ramp(d, -0.4); + let c = ramp(d, 0.5); + + // SAFETY: all 5 slices have length 16, a multiple of 8. + let got = unsafe { micro_4x1(&p0, &p1, &p2, &p3, &c) }; + + let expected = [ + ref_dot(&p0, &c), + ref_dot(&p1, &c), + ref_dot(&p2, &c), + ref_dot(&p3, &c), + ]; + + for (g, e) in got.iter().zip(&expected) { + assert_close(*g, *e, 1e-5); + } + } + + #[test] + fn micro_4x1_d3072() { + let d = 3072; + let p0 = ramp(d, 0.001); + let p1 = ramp(d, -0.001); + let p2 = ramp(d, 0.002); + let p3 = ramp(d, -0.002); + let c = ramp(d, 0.001); + + // SAFETY: all 5 slices have length 3072, a multiple of 8. + let got = unsafe { micro_4x1(&p0, &p1, &p2, &p3, &c) }; + + let expected = [ + ref_dot(&p0, &c), + ref_dot(&p1, &c), + ref_dot(&p2, &c), + ref_dot(&p3, &c), + ]; + + for (g, e) in got.iter().zip(&expected) { + assert_close(*g, *e, 1e-3); + } + } + #[test] fn nearest4_matches_brute_force_even_k() { let d = 8; - let k = 4; + let k = nz!(4); // 4 centroids: axis-aligned unit vectors. - let mut centroids = vec![0.0_f32; k * d]; - for i in 0..k { + let mut centroids = vec![0.0_f32; k.get() * d]; + for i in 0..k.get() { centroids[i * d + i] = 1.0; } @@ -721,10 +732,10 @@ mod tests { #[test] fn nearest4_matches_brute_force_odd_k() { let d = 8; - let k = 3; // odd: exercises the remainder path + let k = nz!(3); // odd: exercises the remainder path - let mut centroids = vec![0.0_f32; k * d]; - for i in 0..k { + let mut centroids = vec![0.0_f32; k.get() * d]; + for i in 0..k.get() { centroids[i * d + i] = 1.0; } @@ -759,7 +770,7 @@ mod tests { // SAFETY: d=8 (multiple of 8), k=1 > 0, centroids has length d, // all point slices have length d. - let got = unsafe { nearest4(&p0, &p1, &p2, &p3, ¢roids, 1, d) }; + let got = unsafe { nearest4(&p0, &p1, &p2, &p3, ¢roids, nz!(1), d) }; assert_eq!(got[0].0, 0); assert_eq!(got[1].0, 0); diff --git a/libs/@local/graph/store/src/embedding/mod.rs b/libs/@local/graph/store/src/embedding/mod.rs index d008b247ce0..c370176aacc 100644 --- a/libs/@local/graph/store/src/embedding/mod.rs +++ b/libs/@local/graph/store/src/embedding/mod.rs @@ -10,4 +10,7 @@ pub mod clustering; pub mod dimension; -pub(crate) mod kernel; +// Hidden from docs: the kernel is an implementation detail, exposed only so +// the `embedding` bench target can measure it in isolation. +#[doc(hidden)] +pub mod kernel; diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index e6c43cdea12..8bf5610f40b 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -576,6 +576,11 @@ pub struct ClusterEntitiesResponse { pub clusters: Vec, /// Entities from the request that had no stored embedding. pub missing_embeddings: Vec, + /// Sum of squared chord distances from every clustered entity to its + /// assigned centroid. Lower is tighter; comparable across runs over the + /// same entities, e.g. to choose a cluster count. `0.0` when nothing was + /// clustered. + pub inertia: f32, } #[derive(Debug, Deserialize)] From 6865e313077a9e0ed31d9d77ce8016a46d6a6018 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:23:27 +0200 Subject: [PATCH 12/38] chore: diagram --- apps/hash-graph/docs/dependency-diagram.mmd | 1 + .../rust/docs/dependency-diagram.mmd | 30 ++++++++++------- .../graph/api/docs/dependency-diagram.mmd | 1 + .../authorization/docs/dependency-diagram.mmd | 30 ++++++++++------- .../embeddings/docs/dependency-diagram.mmd | 24 +++++++++----- .../docs/dependency-diagram.mmd | 32 ++++++++++++------- .../graph/store/docs/dependency-diagram.mmd | 30 ++++++++++------- libs/@local/graph/store/package.json | 1 + .../type-fetcher/docs/dependency-diagram.mmd | 24 +++++++++----- .../graph/types/docs/dependency-diagram.mmd | 30 ++++++++++------- .../validation/docs/dependency-diagram.mmd | 30 ++++++++++------- .../harpc/server/docs/dependency-diagram.mmd | 26 +++++++++------ .../hashql/ast/docs/dependency-diagram.mmd | 1 + .../compiletest/docs/dependency-diagram.mmd | 1 + .../hashql/eval/docs/dependency-diagram.mmd | 1 + .../hashql/hir/docs/dependency-diagram.mmd | 1 + .../hashql/mir/docs/dependency-diagram.mmd | 1 + .../syntax-jexpr/docs/dependency-diagram.mmd | 1 + .../docs/dependency-diagram.mmd | 30 ++++++++++------- .../rust/docs/dependency-diagram.mmd | 32 ++++++++++++------- 20 files changed, 212 insertions(+), 115 deletions(-) diff --git a/apps/hash-graph/docs/dependency-diagram.mmd b/apps/hash-graph/docs/dependency-diagram.mmd index 2c772cf559f..ec8e7669a4d 100644 --- a/apps/hash-graph/docs/dependency-diagram.mmd +++ b/apps/hash-graph/docs/dependency-diagram.mmd @@ -72,6 +72,7 @@ graph TD 10 --> 5 10 --> 14 10 --> 35 + 10 -.-> 37 11 --> 2 12 --> 33 13 --> 10 diff --git a/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd b/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd index d97925ceaff..3cbfa718cbd 100644 --- a/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd +++ b/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd @@ -32,13 +32,17 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[error-stack] - 24[hash-graph-benches] - 25[hash-graph-integration] - 26[hash-graph-test-data] + 23[darwin-kperf] + 24[darwin-kperf-criterion] + 25[darwin-kperf-events] + 26[darwin-kperf-sys] + 27[error-stack] + 28[hash-graph-benches] + 29[hash-graph-integration] + 30[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 26 + 1 -.-> 30 2 -.-> 3 2 --> 15 4 --> 6 @@ -52,15 +56,16 @@ graph TD 8 --> 5 8 --> 11 8 --> 22 + 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 26 - 12 -.-> 26 + 11 -.-> 30 + 12 -.-> 30 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 23 + 15 --> 27 16 -.-> 17 17 --> 18 17 --> 21 @@ -70,6 +75,9 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 24 -.-> 4 - 25 -.-> 7 - 26 --> 8 + 23 --> 25 + 23 --> 26 + 24 --> 23 + 28 -.-> 4 + 29 -.-> 7 + 30 --> 8 diff --git a/libs/@local/graph/api/docs/dependency-diagram.mmd b/libs/@local/graph/api/docs/dependency-diagram.mmd index 3ec3bd7bc71..ec3530323a2 100644 --- a/libs/@local/graph/api/docs/dependency-diagram.mmd +++ b/libs/@local/graph/api/docs/dependency-diagram.mmd @@ -73,6 +73,7 @@ graph TD 10 --> 5 10 --> 14 10 --> 35 + 10 -.-> 37 11 --> 2 12 --> 33 13 --> 10 diff --git a/libs/@local/graph/authorization/docs/dependency-diagram.mmd b/libs/@local/graph/authorization/docs/dependency-diagram.mmd index 9b7424d94ae..002c88588d2 100644 --- a/libs/@local/graph/authorization/docs/dependency-diagram.mmd +++ b/libs/@local/graph/authorization/docs/dependency-diagram.mmd @@ -32,13 +32,17 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[error-stack] - 24[hash-graph-benches] - 25[hash-graph-integration] - 26[hash-graph-test-data] + 23[darwin-kperf] + 24[darwin-kperf-criterion] + 25[darwin-kperf-events] + 26[darwin-kperf-sys] + 27[error-stack] + 28[hash-graph-benches] + 29[hash-graph-integration] + 30[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 26 + 1 -.-> 30 2 -.-> 3 2 --> 15 4 --> 6 @@ -52,15 +56,16 @@ graph TD 8 --> 5 8 --> 11 8 --> 22 + 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 26 - 12 -.-> 26 + 11 -.-> 30 + 12 -.-> 30 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 23 + 15 --> 27 16 -.-> 17 17 --> 18 17 --> 21 @@ -70,6 +75,9 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 24 -.-> 4 - 25 -.-> 7 - 26 --> 8 + 23 --> 25 + 23 --> 26 + 24 --> 23 + 28 -.-> 4 + 29 -.-> 7 + 30 --> 8 diff --git a/libs/@local/graph/embeddings/docs/dependency-diagram.mmd b/libs/@local/graph/embeddings/docs/dependency-diagram.mmd index 1ebad383105..97ab6d97b76 100644 --- a/libs/@local/graph/embeddings/docs/dependency-diagram.mmd +++ b/libs/@local/graph/embeddings/docs/dependency-diagram.mmd @@ -22,12 +22,16 @@ graph TD 10[harpc-types] 11[harpc-wire-protocol] 12[hash-temporal-client] - 13[error-stack] - 14[hash-graph-benches] - 15[hash-graph-test-data] + 13[darwin-kperf] + 14[darwin-kperf-criterion] + 15[darwin-kperf-events] + 16[darwin-kperf-sys] + 17[error-stack] + 18[hash-graph-benches] + 19[hash-graph-test-data] 0 --> 4 1 --> 8 - 1 -.-> 15 + 1 -.-> 19 2 -.-> 3 2 --> 11 4 --> 6 @@ -37,11 +41,15 @@ graph TD 7 --> 5 7 --> 9 7 --> 12 + 7 -.-> 14 8 --> 2 - 9 -.-> 15 + 9 -.-> 19 11 -.-> 10 11 --> 10 - 11 --> 13 + 11 --> 17 12 --> 1 - 14 -.-> 4 - 15 --> 7 + 13 --> 15 + 13 --> 16 + 14 --> 13 + 18 -.-> 4 + 19 --> 7 diff --git a/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd b/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd index 84ed6a1efa0..d26e8f3bc11 100644 --- a/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd +++ b/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd @@ -33,13 +33,17 @@ graph TD 21[hash-status] 22[hash-telemetry] 23[hash-temporal-client] - 24[error-stack] - 25[hash-graph-benches] - 26[hash-graph-integration] - 27[hash-graph-test-data] + 24[darwin-kperf] + 25[darwin-kperf-criterion] + 26[darwin-kperf-events] + 27[darwin-kperf-sys] + 28[error-stack] + 29[hash-graph-benches] + 30[hash-graph-integration] + 31[hash-graph-test-data] 0 --> 4 1 --> 10 - 1 -.-> 27 + 1 -.-> 31 2 -.-> 3 2 --> 14 4 --> 17 @@ -53,12 +57,13 @@ graph TD 9 --> 5 9 --> 11 9 --> 23 + 9 -.-> 25 10 --> 2 - 11 -.-> 27 - 12 -.-> 27 + 11 -.-> 31 + 12 -.-> 31 14 -.-> 13 14 --> 13 - 14 --> 24 + 14 --> 28 15 -.-> 16 16 --> 17 16 --> 20 @@ -67,8 +72,11 @@ graph TD 18 -.-> 16 19 --> 18 20 --> 15 - 22 --> 24 + 22 --> 28 23 --> 1 - 25 -.-> 4 - 26 -.-> 8 - 27 --> 9 + 24 --> 26 + 24 --> 27 + 25 --> 24 + 29 -.-> 4 + 30 -.-> 8 + 31 --> 9 diff --git a/libs/@local/graph/store/docs/dependency-diagram.mmd b/libs/@local/graph/store/docs/dependency-diagram.mmd index 0ba92b32822..c25fc15d926 100644 --- a/libs/@local/graph/store/docs/dependency-diagram.mmd +++ b/libs/@local/graph/store/docs/dependency-diagram.mmd @@ -32,13 +32,17 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[error-stack] - 24[hash-graph-benches] - 25[hash-graph-integration] - 26[hash-graph-test-data] + 23[darwin-kperf] + 24[darwin-kperf-criterion] + 25[darwin-kperf-events] + 26[darwin-kperf-sys] + 27[error-stack] + 28[hash-graph-benches] + 29[hash-graph-integration] + 30[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 26 + 1 -.-> 30 2 -.-> 3 2 --> 15 4 --> 6 @@ -52,15 +56,16 @@ graph TD 8 --> 5 8 --> 11 8 --> 22 + 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 26 - 12 -.-> 26 + 11 -.-> 30 + 12 -.-> 30 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 23 + 15 --> 27 16 -.-> 17 17 --> 18 17 --> 21 @@ -70,6 +75,9 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 24 -.-> 4 - 25 -.-> 7 - 26 --> 8 + 23 --> 25 + 23 --> 26 + 24 --> 23 + 28 -.-> 4 + 29 -.-> 7 + 30 --> 8 diff --git a/libs/@local/graph/store/package.json b/libs/@local/graph/store/package.json index 13a00b5532b..04978d129c6 100644 --- a/libs/@local/graph/store/package.json +++ b/libs/@local/graph/store/package.json @@ -33,6 +33,7 @@ }, "devDependencies": { "@local/tsconfig": "workspace:*", + "@rust/darwin-kperf-criterion": "workspace:*", "@rust/hash-codegen": "workspace:*", "typescript": "5.9.3" } diff --git a/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd b/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd index 5c3c10fd8c7..d1243620b21 100644 --- a/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd +++ b/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd @@ -22,12 +22,16 @@ graph TD 10[harpc-types] 11[harpc-wire-protocol] 12[hash-temporal-client] - 13[error-stack] - 14[hash-graph-benches] - 15[hash-graph-test-data] + 13[darwin-kperf] + 14[darwin-kperf-criterion] + 15[darwin-kperf-events] + 16[darwin-kperf-sys] + 17[error-stack] + 18[hash-graph-benches] + 19[hash-graph-test-data] 0 --> 4 1 --> 7 - 1 -.-> 15 + 1 -.-> 19 2 -.-> 3 2 --> 11 4 --> 8 @@ -35,12 +39,16 @@ graph TD 6 --> 5 6 --> 9 6 --> 12 + 6 -.-> 14 7 --> 2 8 --> 6 - 9 -.-> 15 + 9 -.-> 19 11 -.-> 10 11 --> 10 - 11 --> 13 + 11 --> 17 12 --> 1 - 14 -.-> 4 - 15 --> 6 + 13 --> 15 + 13 --> 16 + 14 --> 13 + 18 -.-> 4 + 19 --> 6 diff --git a/libs/@local/graph/types/docs/dependency-diagram.mmd b/libs/@local/graph/types/docs/dependency-diagram.mmd index 64f359477cc..1b887fb64f8 100644 --- a/libs/@local/graph/types/docs/dependency-diagram.mmd +++ b/libs/@local/graph/types/docs/dependency-diagram.mmd @@ -32,13 +32,17 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[error-stack] - 24[hash-graph-benches] - 25[hash-graph-integration] - 26[hash-graph-test-data] + 23[darwin-kperf] + 24[darwin-kperf-criterion] + 25[darwin-kperf-events] + 26[darwin-kperf-sys] + 27[error-stack] + 28[hash-graph-benches] + 29[hash-graph-integration] + 30[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 26 + 1 -.-> 30 2 -.-> 3 2 --> 15 4 --> 6 @@ -52,15 +56,16 @@ graph TD 8 --> 5 8 --> 11 8 --> 22 + 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 26 - 12 -.-> 26 + 11 -.-> 30 + 12 -.-> 30 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 23 + 15 --> 27 16 -.-> 17 17 --> 18 17 --> 21 @@ -70,6 +75,9 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 24 -.-> 4 - 25 -.-> 7 - 26 --> 8 + 23 --> 25 + 23 --> 26 + 24 --> 23 + 28 -.-> 4 + 29 -.-> 7 + 30 --> 8 diff --git a/libs/@local/graph/validation/docs/dependency-diagram.mmd b/libs/@local/graph/validation/docs/dependency-diagram.mmd index a84b3c63ba9..4dd6fa9d4c7 100644 --- a/libs/@local/graph/validation/docs/dependency-diagram.mmd +++ b/libs/@local/graph/validation/docs/dependency-diagram.mmd @@ -29,13 +29,17 @@ graph TD 17[hashql-mir] 18[hashql-syntax-jexpr] 19[hash-temporal-client] - 20[error-stack] - 21[hash-graph-benches] - 22[hash-graph-integration] - 23[hash-graph-test-data] + 20[darwin-kperf] + 21[darwin-kperf-criterion] + 22[darwin-kperf-events] + 23[darwin-kperf-sys] + 24[error-stack] + 25[hash-graph-benches] + 26[hash-graph-integration] + 27[hash-graph-test-data] 0 --> 4 1 --> 8 - 1 -.-> 23 + 1 -.-> 27 2 -.-> 3 2 --> 12 4 --> 15 @@ -45,12 +49,13 @@ graph TD 7 --> 5 7 --> 9 7 --> 19 + 7 -.-> 21 8 --> 2 - 9 -.-> 23 - 10 -.-> 23 + 9 -.-> 27 + 10 -.-> 27 12 -.-> 11 12 --> 11 - 12 --> 20 + 12 --> 24 13 -.-> 14 14 --> 15 14 --> 18 @@ -60,6 +65,9 @@ graph TD 17 --> 16 18 --> 13 19 --> 1 - 21 -.-> 4 - 22 -.-> 6 - 23 --> 7 + 20 --> 22 + 20 --> 23 + 21 --> 20 + 25 -.-> 4 + 26 -.-> 6 + 27 --> 7 diff --git a/libs/@local/harpc/server/docs/dependency-diagram.mmd b/libs/@local/harpc/server/docs/dependency-diagram.mmd index c55bd6573e3..0f088c57070 100644 --- a/libs/@local/harpc/server/docs/dependency-diagram.mmd +++ b/libs/@local/harpc/server/docs/dependency-diagram.mmd @@ -27,12 +27,16 @@ graph TD 15[harpc-types] 16[harpc-wire-protocol] 17[hash-temporal-client] - 18[error-stack] - 19[hash-graph-benches] - 20[hash-graph-test-data] + 18[darwin-kperf] + 19[darwin-kperf-criterion] + 20[darwin-kperf-events] + 21[darwin-kperf-sys] + 22[error-stack] + 23[hash-graph-benches] + 24[hash-graph-test-data] 0 --> 4 1 --> 7 - 1 -.-> 20 + 1 -.-> 24 2 -.-> 3 2 --> 16 4 --> 6 @@ -41,11 +45,12 @@ graph TD 6 --> 5 6 --> 8 6 --> 17 + 6 -.-> 19 7 --> 2 - 8 -.-> 20 + 8 -.-> 24 9 --> 13 10 --> 15 - 10 --> 18 + 10 --> 22 11 --> 2 11 -.-> 10 11 --> 10 @@ -56,7 +61,10 @@ graph TD 14 --> 11 16 -.-> 15 16 --> 15 - 16 --> 18 + 16 --> 22 17 --> 1 - 19 -.-> 4 - 20 --> 6 + 18 --> 20 + 18 --> 21 + 19 --> 18 + 23 -.-> 4 + 24 --> 6 diff --git a/libs/@local/hashql/ast/docs/dependency-diagram.mmd b/libs/@local/hashql/ast/docs/dependency-diagram.mmd index 1c4dfe4d89f..1d809d023df 100644 --- a/libs/@local/hashql/ast/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/ast/docs/dependency-diagram.mmd @@ -59,6 +59,7 @@ graph TD 9 --> 5 9 --> 11 9 --> 26 + 9 -.-> 28 10 --> 2 11 -.-> 33 12 -.-> 33 diff --git a/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd b/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd index 4d2d74f3e20..7fa03c59383 100644 --- a/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd @@ -59,6 +59,7 @@ graph TD 9 --> 5 9 --> 11 9 --> 26 + 9 -.-> 28 10 --> 2 11 -.-> 33 12 -.-> 33 diff --git a/libs/@local/hashql/eval/docs/dependency-diagram.mmd b/libs/@local/hashql/eval/docs/dependency-diagram.mmd index 09fbd7269eb..ba720211915 100644 --- a/libs/@local/hashql/eval/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/eval/docs/dependency-diagram.mmd @@ -59,6 +59,7 @@ graph TD 9 --> 5 9 --> 11 9 --> 26 + 9 -.-> 28 10 --> 2 11 -.-> 33 12 -.-> 33 diff --git a/libs/@local/hashql/hir/docs/dependency-diagram.mmd b/libs/@local/hashql/hir/docs/dependency-diagram.mmd index 979a0dc9dc4..31c175f2525 100644 --- a/libs/@local/hashql/hir/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/hir/docs/dependency-diagram.mmd @@ -59,6 +59,7 @@ graph TD 9 --> 5 9 --> 11 9 --> 26 + 9 -.-> 28 10 --> 2 11 -.-> 33 12 -.-> 33 diff --git a/libs/@local/hashql/mir/docs/dependency-diagram.mmd b/libs/@local/hashql/mir/docs/dependency-diagram.mmd index 58b887e6182..351c3ebd1f6 100644 --- a/libs/@local/hashql/mir/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/mir/docs/dependency-diagram.mmd @@ -59,6 +59,7 @@ graph TD 9 --> 5 9 --> 11 9 --> 26 + 9 -.-> 28 10 --> 2 11 -.-> 33 12 -.-> 33 diff --git a/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd b/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd index d4de1967175..1abbd2f75d5 100644 --- a/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd @@ -59,6 +59,7 @@ graph TD 9 --> 5 9 --> 11 9 --> 26 + 9 -.-> 28 10 --> 2 11 -.-> 33 12 -.-> 33 diff --git a/libs/@local/temporal-client/docs/dependency-diagram.mmd b/libs/@local/temporal-client/docs/dependency-diagram.mmd index 65bcb84bf66..1152c9ec8ed 100644 --- a/libs/@local/temporal-client/docs/dependency-diagram.mmd +++ b/libs/@local/temporal-client/docs/dependency-diagram.mmd @@ -32,13 +32,17 @@ graph TD 21[hashql-syntax-jexpr] 22[hash-temporal-client] class 22 root - 23[error-stack] - 24[hash-graph-benches] - 25[hash-graph-integration] - 26[hash-graph-test-data] + 23[darwin-kperf] + 24[darwin-kperf-criterion] + 25[darwin-kperf-events] + 26[darwin-kperf-sys] + 27[error-stack] + 28[hash-graph-benches] + 29[hash-graph-integration] + 30[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 26 + 1 -.-> 30 2 -.-> 3 2 --> 15 4 --> 6 @@ -52,15 +56,16 @@ graph TD 8 --> 5 8 --> 11 8 --> 22 + 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 26 - 12 -.-> 26 + 11 -.-> 30 + 12 -.-> 30 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 23 + 15 --> 27 16 -.-> 17 17 --> 18 17 --> 21 @@ -70,6 +75,9 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 24 -.-> 4 - 25 -.-> 7 - 26 --> 8 + 23 --> 25 + 23 --> 26 + 24 --> 23 + 28 -.-> 4 + 29 -.-> 7 + 30 --> 8 diff --git a/tests/graph/test-data/rust/docs/dependency-diagram.mmd b/tests/graph/test-data/rust/docs/dependency-diagram.mmd index 8418e7f6ed1..fc41eb6288a 100644 --- a/tests/graph/test-data/rust/docs/dependency-diagram.mmd +++ b/tests/graph/test-data/rust/docs/dependency-diagram.mmd @@ -31,14 +31,18 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[error-stack] - 24[hash-graph-benches] - 25[hash-graph-integration] - 26[hash-graph-test-data] - class 26 root + 23[darwin-kperf] + 24[darwin-kperf-criterion] + 25[darwin-kperf-events] + 26[darwin-kperf-sys] + 27[error-stack] + 28[hash-graph-benches] + 29[hash-graph-integration] + 30[hash-graph-test-data] + class 30 root 0 --> 4 1 --> 9 - 1 -.-> 26 + 1 -.-> 30 2 -.-> 3 2 --> 15 4 --> 6 @@ -52,15 +56,16 @@ graph TD 8 --> 5 8 --> 11 8 --> 22 + 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 26 - 12 -.-> 26 + 11 -.-> 30 + 12 -.-> 30 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 23 + 15 --> 27 16 -.-> 17 17 --> 18 17 --> 21 @@ -70,6 +75,9 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 24 -.-> 4 - 25 -.-> 7 - 26 --> 8 + 23 --> 25 + 23 --> 26 + 24 --> 23 + 28 -.-> 4 + 29 -.-> 7 + 30 --> 8 From fa9b35a2f6894ea4f05b2c8bbf57b668efbdef47 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:23:45 +0200 Subject: [PATCH 13/38] chore: lockfile --- yarn.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/yarn.lock b/yarn.lock index 8e112e18814..cb4ccc74b26 100644 --- a/yarn.lock +++ b/yarn.lock @@ -13215,6 +13215,7 @@ __metadata: dependencies: "@blockprotocol/type-system-rs": "workspace:*" "@local/tsconfig": "workspace:*" + "@rust/darwin-kperf-criterion": "workspace:*" "@rust/error-stack": "workspace:*" "@rust/hash-codec": "workspace:*" "@rust/hash-codegen": "workspace:*" From 2bfd8badad73eaf5a9f4707f8b31ebbccdf3d0f4 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Fri, 3 Jul 2026 18:20:37 +0200 Subject: [PATCH 14/38] chore: fix schema --- libs/@local/graph/api/openapi/openapi.json | 3 ++- libs/@local/graph/store/src/entity/store.rs | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index a93edfd1e59..2839c9c7d79 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -3809,7 +3809,8 @@ "description": "Embedding dimension after matryoshka truncation. Must be a positive\nmultiple of 8; values above 3072 are rejected. Defaults to 256.", "default": 256, "example": 256, - "minimum": 1 + "multipleOf": 8, + "minimum": 8 }, "entityIds": { "type": "array", diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 8bf5610f40b..745333971a5 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -539,7 +539,7 @@ pub struct ClusterEntitiesParams { /// Embedding dimension after matryoshka truncation. Must be a positive /// multiple of 8; values above 3072 are rejected. Defaults to 256. #[serde(default = "ClusterEntitiesParams::default_dimension")] - #[cfg_attr(feature = "utoipa", schema(value_type = u16, minimum = 1, default = 256, example = 256))] + #[cfg_attr(feature = "utoipa", schema(value_type = u16, minimum = 8, multiple_of = 8, default = 256, example = 256))] pub dimension: NonZero, /// Seed for the random number generator used in clustering. From 3cc5897e20d8d88f0b4135f1785bc53ba37511b0 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 09:57:06 +0200 Subject: [PATCH 15/38] feat: add scripts --- libs/@local/graph/store/package.json | 3 +++ 1 file changed, 3 insertions(+) diff --git a/libs/@local/graph/store/package.json b/libs/@local/graph/store/package.json index 04978d129c6..2fb5ede493d 100644 --- a/libs/@local/graph/store/package.json +++ b/libs/@local/graph/store/package.json @@ -15,11 +15,14 @@ "./types": "./types/index.snap.js" }, "scripts": { + "build:codspeed": "cargo codspeed build -p hash-graph-store", "build:types": "INSTA_UPDATE=always mise exec --env dev cargo:cargo-insta -- cargo-insta test --features codegen --test codegen", "doc:dependency-diagram": "cargo run -p hash-repo-chores -- dependency-diagram --output docs/dependency-diagram.mmd --root hash-graph-store --root-deps-and-dependents --link-mode non-roots --include-dev-deps --include-build-deps --logging-console-level info", "fix:clippy": "just clippy --fix", "lint:clippy": "just clippy", "lint:tsc": "tsc --noEmit", + "test:codspeed": "cargo codspeed run -p hash-graph-store", + "test:miri": "cargo miri nextest run -- embedding::kernel embedding::clustering::tests::nearest_centroid_argmax_independent_of_inv_norm", "test:unit": "mise run test:unit @rust/hash-graph-store" }, "dependencies": { From 7a8ddad5ffe315c8b8ffe1326c0057e0377cd73e Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:14:55 +0200 Subject: [PATCH 16/38] fix: warm up repository on CI --- .github/workflows/codspeed.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 8a824c166ab..867aced9ae0 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -70,6 +70,9 @@ jobs: with: scope: ${{ matrix.name }} + - name: Warm up repository + uses: ./.github/actions/warm-up-repo + - name: Build the benchmark target run: turbo run build:codspeed --filter=${{ matrix.name }} From 336dd2f17c677df889f02009de376eb8689e2916 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:30:10 +0200 Subject: [PATCH 17/38] fix: docs --- libs/@local/graph/api/openapi/openapi.json | 2 +- libs/@local/graph/store/src/entity/store.rs | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index 2839c9c7d79..f32b0eb637b 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -4638,7 +4638,7 @@ "type": "number", "format": "float" }, - "description": "Unit-normalized centroid with length equal to the requested dimension." + "description": "Centroid with length equal to the requested dimension.\n\nTypically unit-normalized, but may be the all-zero vector if all assigned points have zero\nnorm." }, "clusterId": { "type": "integer", diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 745333971a5..14ecc7a08ee 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -562,7 +562,10 @@ pub struct EntityCluster { /// Index in `0..cluster_count`. pub cluster_id: u16, pub entity_ids: Vec, - /// Unit-normalized centroid with length equal to the requested dimension. + /// Centroid with length equal to the requested dimension. + /// + /// Typically unit-normalized, but may be the all-zero vector if all assigned points have zero + /// norm. pub centroid: Vec, } From 4685963995cfcb253f2b00f53d2f252775675d2c Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 13:04:38 +0200 Subject: [PATCH 18/38] chore: bound the threads inside of benchmarking + thread pool limit --- Cargo.lock | 1 + apps/hash-graph/Cargo.toml | 1 + apps/hash-graph/src/subcommand/mod.rs | 13 ++++++++++++- libs/@local/graph/store/benches/embedding.rs | 5 +++++ 4 files changed, 19 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 8c32c3395bf..3045f2d5588 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3527,6 +3527,7 @@ dependencies = [ "jsonwebtoken", "mimalloc", "multiaddr", + "rayon", "regex", "reqwest", "simple-mermaid", diff --git a/apps/hash-graph/Cargo.toml b/apps/hash-graph/Cargo.toml index 0dcb10e4746..02b4f58de69 100644 --- a/apps/hash-graph/Cargo.toml +++ b/apps/hash-graph/Cargo.toml @@ -34,6 +34,7 @@ futures = { workspace = true } jsonwebtoken = { workspace = true } mimalloc = { workspace = true } multiaddr = { workspace = true } +rayon = { workspace = true } regex = { workspace = true } reqwest = { workspace = true, features = ["rustls"] } simple-mermaid = { workspace = true } diff --git a/apps/hash-graph/src/subcommand/mod.rs b/apps/hash-graph/src/subcommand/mod.rs index 08931291972..e6e7ca8e334 100644 --- a/apps/hash-graph/src/subcommand/mod.rs +++ b/apps/hash-graph/src/subcommand/mod.rs @@ -7,7 +7,7 @@ mod snapshot; mod type_fetcher; use core::time::Duration; -use std::time::Instant; +use std::{thread::available_parallelism, time::Instant}; use clap::Parser; use error_stack::{Report, ensure}; @@ -141,11 +141,22 @@ pub enum Subcommand { ReindexCache(Box), } +#[expect(clippy::integer_division, clippy::integer_division_remainder_used)] fn block_on( future: impl Future>>, service_name: &'static str, tracing_config: TracingConfig, ) -> Result<(), Report> { + rayon::ThreadPoolBuilder::new() + .num_threads( + available_parallelism() + .map_or(1, |cores| cores.get() / 2) + .max(1), + ) + .thread_name(|index| format!("rayon-{index}")) + .build_global() + .expect("rayon pool should be initialized exactly once"); + tokio::runtime::Builder::new_multi_thread() .enable_all() .build() diff --git a/libs/@local/graph/store/benches/embedding.rs b/libs/@local/graph/store/benches/embedding.rs index 13094b45589..2bd5b53692b 100644 --- a/libs/@local/graph/store/benches/embedding.rs +++ b/libs/@local/graph/store/benches/embedding.rs @@ -179,6 +179,11 @@ fn bench_nearest4(criterion: &mut Criterion) { } fn bench_cluster(criterion: &mut Criterion) { + rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build_global() + .expect("should be built exactly once"); + let mut group = criterion.benchmark_group("embedding/cluster"); group.sample_size(10); From fa9c9da1eb837a148644b8dbcfdc94893296d44f Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:40:46 +0200 Subject: [PATCH 19/38] feat: limit request + remove clustering unsafe --- .../store/postgres/knowledge/entity/mod.rs | 18 +++++++++--- .../graph/store/src/embedding/clustering.rs | 29 +++++-------------- libs/@local/graph/store/src/embedding/mod.rs | 4 +-- libs/@local/graph/store/src/entity/store.rs | 15 ++++++---- libs/@local/graph/store/src/error.rs | 2 ++ 5 files changed, 35 insertions(+), 33 deletions(-) diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 4faecf6befb..d69d276cfaf 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -2599,7 +2599,7 @@ where Ok(permitted_ids) } - #[expect(clippy::too_many_lines, clippy::cast_possible_truncation)] + #[expect(clippy::too_many_lines)] #[tracing::instrument(skip(self, params))] async fn cluster_entities( &self, @@ -2607,7 +2607,8 @@ where params: ClusterEntitiesParams, ) -> Result> { const { assert!(Embedding::DIM <= u16::MAX as usize) }; - const STORED_DIM: u16 = Embedding::DIM as u16; + const MAX_ALLOWED_DIM: u16 = 512; + const MAX_ALLOWED_K: u16 = 64; let dimension = Dimension::new(params.dimension.get()).ok_or_else(|| { Report::new(ClusterError::InvalidDimension { @@ -2616,13 +2617,22 @@ where .attach(StatusCode::InvalidArgument) })?; - if dimension.get() > STORED_DIM { + if dimension.get() > MAX_ALLOWED_DIM { return Err(Report::new(ClusterError::DimensionTooLarge { dimension: dimension.value(), - max: STORED_DIM, + max: MAX_ALLOWED_DIM, + }) + .attach(StatusCode::InvalidArgument)); + } + + if params.cluster_count > MAX_ALLOWED_K { + return Err(Report::new(ClusterError::KTooLarge { + count: params.cluster_count, + max: MAX_ALLOWED_K, }) .attach(StatusCode::InvalidArgument)); } + let truncated_dim = usize::from(dimension.get()); // Filter to entities the actor is allowed to view. diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index 0ff02308ee7..03a0574fe9a 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -100,13 +100,8 @@ pub struct Clustering { impl Clustering { /// Allocates a zeroed clustering for `k` centroids over `n` points. fn new(k: u16, n: usize, d: Dimension) -> Self { - // SAFETY: All-zero bits are valid for `f32` (IEEE 754 positive zero) - // and for `u16` (the integer 0). `Box::new_zeroed_slice` allocates - // zeroed memory of the correct layout, so `assume_init` is sound. - let centroids: Box<[f32]> = - unsafe { Box::new_zeroed_slice((k as usize) * (d.get() as usize)).assume_init() }; - // SAFETY: All-zero bits are valid for `u16` (the integer 0). - let labels: Box<[u16]> = unsafe { Box::new_zeroed_slice(n).assume_init() }; + let centroids: Box<[f32]> = vec![0.0; (k as usize) * (d.get() as usize)].into_boxed_slice(); + let labels: Box<[u16]> = vec![0; n].into_boxed_slice(); Self { centroids, @@ -329,20 +324,12 @@ struct Restart { impl Restart { fn new(k: NonZero, m: usize, d: usize) -> Self { - // SAFETY: all-zero bits are valid for `f32` (IEEE 754 +0.0), `usize` (0), `u16` (0), and - // `bool` (false). `Box::new_zeroed_slice` allocates zeroed memory of the correct - // layout for each type, so `assume_init` is sound in every case. - let centroids = unsafe { Box::<[f32]>::new_zeroed_slice(k.get() * d).assume_init() }; - // SAFETY: see above - let sums = unsafe { Box::<[f32]>::new_zeroed_slice(k.get() * d).assume_init() }; - // SAFETY: see above - let counts = unsafe { Box::<[usize]>::new_zeroed_slice(k.get()).assume_init() }; - // SAFETY: see above - let labels = unsafe { Box::<[u16]>::new_zeroed_slice(m).assume_init() }; - // SAFETY: see above - let point_distances = unsafe { Box::<[f32]>::new_zeroed_slice(m).assume_init() }; - // SAFETY: see above - let selected = unsafe { Box::<[bool]>::new_zeroed_slice(m).assume_init() }; + let centroids: Box<[f32]> = vec![0.0; k.get() * d].into_boxed_slice(); + let sums: Box<[f32]> = vec![0.0; k.get() * d].into_boxed_slice(); + let counts: Box<[usize]> = vec![0; k.get()].into_boxed_slice(); + let labels: Box<[u16]> = vec![0; m].into_boxed_slice(); + let point_distances: Box<[f32]> = vec![0.0; m].into_boxed_slice(); + let selected: Box<[bool]> = vec![false; m].into_boxed_slice(); Self { k, diff --git a/libs/@local/graph/store/src/embedding/mod.rs b/libs/@local/graph/store/src/embedding/mod.rs index c370176aacc..d184500214f 100644 --- a/libs/@local/graph/store/src/embedding/mod.rs +++ b/libs/@local/graph/store/src/embedding/mod.rs @@ -4,8 +4,8 @@ clippy::float_arithmetic, clippy::min_ident_chars, clippy::many_single_char_names, - reason = "embedding module is under active development. Single-char idents (k, n, m, d, x) \ - are standard mathematical notation for clustering." + reason = "Single-char idents (k, n, m, d, x) are standard mathematical notation for \ + clustering." )] pub mod clustering; diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 14ecc7a08ee..59c5532b7c8 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -533,13 +533,16 @@ impl PatchEntityParams { #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct ClusterEntitiesParams { pub entity_ids: Vec, - /// Desired number of clusters. Clamped to the number of entities with - /// embeddings when that is smaller. + /// Desired number of clusters. + /// + /// Clamped to the number of entities with embeddings when that is smaller. + #[cfg_attr(feature = "utoipa", schema(minimum = 0, maximum = 64))] pub cluster_count: u16, - /// Embedding dimension after matryoshka truncation. Must be a positive - /// multiple of 8; values above 3072 are rejected. Defaults to 256. + /// Embedding dimension after matryoshka truncation. + /// + /// Must be a positive multiple of 8; values above 512 are rejected. Defaults to 256. #[serde(default = "ClusterEntitiesParams::default_dimension")] - #[cfg_attr(feature = "utoipa", schema(value_type = u16, minimum = 8, multiple_of = 8, default = 256, example = 256))] + #[cfg_attr(feature = "utoipa", schema(value_type = u16, minimum = 8, maximum = 512, multiple_of = 8, default = 256, example = 256))] pub dimension: NonZero, /// Seed for the random number generator used in clustering. @@ -559,7 +562,7 @@ impl ClusterEntitiesParams { #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))] #[serde(rename_all = "camelCase")] pub struct EntityCluster { - /// Index in `0..cluster_count`. + /// Index in `0..min(cluster_count, n)`. pub cluster_id: u16, pub entity_ids: Vec, /// Centroid with length equal to the requested dimension. diff --git a/libs/@local/graph/store/src/error.rs b/libs/@local/graph/store/src/error.rs index 3f1b684167e..35c35971d67 100644 --- a/libs/@local/graph/store/src/error.rs +++ b/libs/@local/graph/store/src/error.rs @@ -78,6 +78,8 @@ pub enum ClusterError { InvalidDimension { dimension: NonZero }, #[display("dimension {dimension} exceeds stored embedding dimension {max}")] DimensionTooLarge { dimension: NonZero, max: u16 }, + #[display("cluster count {count} exceeds maximum allowed {max}")] + KTooLarge { count: u16, max: u16 }, #[display("embedding query failed")] Store, } From 505d16aa91cc7662a0dd807a4dfd5c97b0200aeb Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:40:44 +0200 Subject: [PATCH 20/38] fix: wording --- libs/@local/graph/store/src/entity/store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 59c5532b7c8..cfd95df07c3 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -580,7 +580,7 @@ pub struct ClusterEntitiesResponse { /// One entry per non-empty cluster. Empty clusters (no points assigned) /// are omitted. pub clusters: Vec, - /// Entities from the request that had no stored embedding. + /// Entities from the request that had no stored embedding or that do not exist. pub missing_embeddings: Vec, /// Sum of squared chord distances from every clustered entity to its /// assigned centroid. Lower is tighter; comparable across runs over the From 7ebfe44837bd44762a091bffa2d42649c6a676fe Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:41:21 +0200 Subject: [PATCH 21/38] chore: regen openapi --- libs/@local/graph/api/openapi/openapi.json | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index f32b0eb637b..332ce2efcaa 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -3800,16 +3800,18 @@ "clusterCount": { "type": "integer", "format": "int32", - "description": "Desired number of clusters. Clamped to the number of entities with\nembeddings when that is smaller.", + "description": "Desired number of clusters.\n\nClamped to the number of entities with embeddings when that is smaller.", + "maximum": 64, "minimum": 0 }, "dimension": { "type": "integer", "format": "int32", - "description": "Embedding dimension after matryoshka truncation. Must be a positive\nmultiple of 8; values above 3072 are rejected. Defaults to 256.", + "description": "Embedding dimension after matryoshka truncation.\n\nMust be a positive multiple of 8; values above 512 are rejected. Defaults to 256.", "default": 256, "example": 256, "multipleOf": 8, + "maximum": 512, "minimum": 8 }, "entityIds": { @@ -3854,7 +3856,7 @@ "items": { "$ref": "#/components/schemas/EntityId" }, - "description": "Entities from the request that had no stored embedding." + "description": "Entities from the request that had no stored embedding or that do not exist." } } }, @@ -4643,7 +4645,7 @@ "clusterId": { "type": "integer", "format": "int32", - "description": "Index in `0..cluster_count`.", + "description": "Index in `0..min(cluster_count, n)`.", "minimum": 0 }, "entityIds": { From 60603e4d7f628682cf64723c93157b2d4c0a0c96 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 18:03:09 +0200 Subject: [PATCH 22/38] chore: docs --- libs/@local/graph/store/src/embedding/clustering.rs | 4 ++++ libs/@local/graph/store/src/entity/store.rs | 2 +- libs/@local/graph/store/src/error.rs | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index 03a0574fe9a..5db1bcb979b 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -635,6 +635,10 @@ unsafe fn accumulate_clusters( // kernel call it precedes. debug_assert!(inv_norms.is_none_or(|norms| norms.len() == labels.len())); + // TODO(PERF): This isn't as performant as it could be, in theory we could first create a + // histogram, and then use that to dispatch over it. A potential optimization opportunity if + // ever required. + sums.par_chunks_exact_mut(d) .zip(counts.par_iter_mut()) .enumerate() diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index cfd95df07c3..063cce3a7ec 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -989,7 +989,7 @@ pub trait EntityStore { /// /// Returns [`ClusterError::InvalidDimension`] if the dimension is not a /// positive multiple of 8, [`ClusterError::DimensionTooLarge`] if it - /// exceeds the stored embedding width, or [`ClusterError::Store`] if the + /// exceeds the maximum allowed dimension, or [`ClusterError::Store`] if the /// embedding query fails. fn cluster_entities( &self, diff --git a/libs/@local/graph/store/src/error.rs b/libs/@local/graph/store/src/error.rs index 35c35971d67..546e8db5e09 100644 --- a/libs/@local/graph/store/src/error.rs +++ b/libs/@local/graph/store/src/error.rs @@ -76,7 +76,7 @@ impl Error for CheckPermissionError {} pub enum ClusterError { #[display("dimension {dimension} is not a positive multiple of 8")] InvalidDimension { dimension: NonZero }, - #[display("dimension {dimension} exceeds stored embedding dimension {max}")] + #[display("dimension {dimension} exceeds maximum allowed dimension {max}")] DimensionTooLarge { dimension: NonZero, max: u16 }, #[display("cluster count {count} exceeds maximum allowed {max}")] KTooLarge { count: u16, max: u16 }, From 7c1dcbe8e5c51839c3f831ed254d8e17bd5ab60c Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:31:11 +0200 Subject: [PATCH 23/38] feat: change the accumulate clusters method --- libs/@local/graph/store/benches/embedding.rs | 14 +- .../graph/store/src/embedding/clustering.rs | 445 +++++++++++++----- .../graph/store/src/embedding/kernel.rs | 29 +- 3 files changed, 368 insertions(+), 120 deletions(-) diff --git a/libs/@local/graph/store/benches/embedding.rs b/libs/@local/graph/store/benches/embedding.rs index 2bd5b53692b..d5dc2268cc3 100644 --- a/libs/@local/graph/store/benches/embedding.rs +++ b/libs/@local/graph/store/benches/embedding.rs @@ -145,14 +145,14 @@ fn bench_nearest4(criterion: &mut Criterion) { // k = 15 exercises the odd-k remainder path. for &(d, k) in &[ - (256, nz!(15)), - (256, nz!(16)), - (256, nz!(64)), - (1536, nz!(16)), - (3072, nz!(16)), + (nz!(256), nz!(15)), + (nz!(256), nz!(16)), + (nz!(256), nz!(64)), + (nz!(1536), nz!(16)), + (nz!(3072), nz!(16)), ] { - let points: Vec> = (0..4).map(|seed| random_vec(d, 30 + seed)).collect(); - let centroids = random_vec(k.get() * d, 40); + let points: Vec> = (0..4).map(|seed| random_vec(d.get(), 30 + seed)).collect(); + let centroids = random_vec(k.get() * d.get(), 40); group.bench_with_input( BenchmarkId::new(format!("d{d}"), k), diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index 5db1bcb979b..3b92fb711cf 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -201,25 +201,25 @@ pub(crate) unsafe fn nearest_centroid( point_inv_norm: f32, centroids: &[f32], k: NonZero, - d: usize, + d: NonZero, ) -> (u16, f32) { - debug_assert_eq!(point.len(), d); - debug_assert_eq!(centroids.len(), k.get() * d); + debug_assert_eq!(point.len(), d.get()); + debug_assert_eq!(centroids.len(), k.get() * d.get()); // SAFETY: the caller guarantees these preconditions. The hints let the // compiler elide bounds checks on the centroid slicing inside the loop. unsafe { - core::hint::assert_unchecked(point.len() == d); - core::hint::assert_unchecked(centroids.len() == k.get() * d); - core::hint::assert_unchecked(d.is_multiple_of(8)); + core::hint::assert_unchecked(point.len() == d.get()); + core::hint::assert_unchecked(centroids.len() == k.get() * d.get()); + core::hint::assert_unchecked(d.get().is_multiple_of(8)); } let mut best = 0; let mut best_dot = f32::NEG_INFINITY; for cluster in 0..k.get() { - let start = cluster * d; - let centroid = ¢roids[start..start + d]; + let start = cluster * d.get(); + let centroid = ¢roids[start..start + d.get()]; // SAFETY: `point` and `centroid` both have length `D`, and `D` is a // multiple of 8 (guaranteed by Dimension). @@ -251,7 +251,7 @@ pub(crate) unsafe fn nearest_centroid( /// * `d` is a multiple of 8 unsafe fn lloyd_assign( k: NonZero, - d: usize, + d: NonZero, centroids: &[f32], points: &[f32], inv_norms: &[f32], @@ -263,18 +263,18 @@ unsafe fn lloyd_assign( // SAFETY: the caller guarantees the length relations; the hints let the // compiler elide bounds checks in the tiled loop below. unsafe { - core::hint::assert_unchecked(points.len() == count * d); + core::hint::assert_unchecked(points.len() == count * d.get()); core::hint::assert_unchecked(inv_norms.len() == count); core::hint::assert_unchecked(distances.len() == count); - core::hint::assert_unchecked(d.is_multiple_of(8)); + core::hint::assert_unchecked(d.get().is_multiple_of(8)); } let mut i = 0; while i + 4 <= count { - let p0 = &points[i * d..i * d + d]; - let p1 = &points[(i + 1) * d..(i + 1) * d + d]; - let p2 = &points[(i + 2) * d..(i + 2) * d + d]; - let p3 = &points[(i + 3) * d..(i + 3) * d + d]; + let p0 = &points[i * d.get()..i * d.get() + d.get()]; + let p1 = &points[(i + 1) * d.get()..(i + 1) * d.get() + d.get()]; + let p2 = &points[(i + 2) * d.get()..(i + 2) * d.get() + d.get()]; + let p3 = &points[(i + 3) * d.get()..(i + 3) * d.get() + d.get()]; // SAFETY: each point length d, centroids length k*d, // k > 0, d a multiple of 8 (guaranteed by Dimension). @@ -289,7 +289,7 @@ unsafe fn lloyd_assign( } while i < count { - let point = &points[i * d..i * d + d]; + let point = &points[i * d.get()..i * d.get() + d.get()]; // SAFETY: point length d, centroids length k*d, k > 0, d mult of 8. let (label, distance) = unsafe { nearest_centroid(point, inv_norms[i], centroids, k, d) }; @@ -306,7 +306,7 @@ unsafe fn lloyd_assign( struct Restart { k: NonZero, m: usize, - d: usize, + d: NonZero, /// Centroids for this restart, `k * d` elements. centroids: Box<[f32]>, @@ -320,16 +320,24 @@ struct Restart { point_distances: Box<[f32]>, /// Tracks which sample points have been selected as seeds. selected: Box<[bool]>, + /// Point indices grouped by cluster, `m` elements; grouping scratch for + /// [`accumulate_clusters`]. + order: Box<[usize]>, + /// Bucket cursors/boundaries, `k + 1` elements; grouping scratch for + /// [`accumulate_clusters`]. + bounds: Box<[usize]>, } impl Restart { - fn new(k: NonZero, m: usize, d: usize) -> Self { - let centroids: Box<[f32]> = vec![0.0; k.get() * d].into_boxed_slice(); - let sums: Box<[f32]> = vec![0.0; k.get() * d].into_boxed_slice(); + fn new(k: NonZero, m: usize, d: NonZero) -> Self { + let centroids: Box<[f32]> = vec![0.0; k.get() * d.get()].into_boxed_slice(); + let sums: Box<[f32]> = vec![0.0; k.get() * d.get()].into_boxed_slice(); let counts: Box<[usize]> = vec![0; k.get()].into_boxed_slice(); let labels: Box<[u16]> = vec![0; m].into_boxed_slice(); let point_distances: Box<[f32]> = vec![0.0; m].into_boxed_slice(); let selected: Box<[bool]> = vec![false; m].into_boxed_slice(); + let order: Box<[usize]> = vec![0; m].into_boxed_slice(); + let bounds: Box<[usize]> = vec![0; k.get() + 1].into_boxed_slice(); Self { k, @@ -341,6 +349,8 @@ impl Restart { labels, point_distances, selected, + order, + bounds, } } @@ -374,15 +384,15 @@ impl Restart { let mut point = rng.random_range(0..m); for cluster in 0..k.get() { - let centroid_start = cluster * d; - let point_start = point * d; + let centroid_start = cluster * d.get(); + let point_start = point * d.get(); - self.centroids[centroid_start..centroid_start + d] - .copy_from_slice(&sample[point_start..point_start + d]); + self.centroids[centroid_start..centroid_start + d.get()] + .copy_from_slice(&sample[point_start..point_start + d.get()]); // SAFETY: centroid rows have length `d`, and `d` is a multiple of 8. unsafe { - kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); + kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d.get()]); } self.selected[point] = true; @@ -393,12 +403,12 @@ impl Restart { break; } - let centroid = &self.centroids[centroid_start..centroid_start + d]; + let centroid = &self.centroids[centroid_start..centroid_start + d.get()]; // Per-element writes only, so the pass is deterministic under // rayon; the D² total is summed sequentially below. sample - .par_chunks_exact(d) + .par_chunks_exact(d.get()) .zip(sample_inv_norms.par_iter()) .zip(self.point_distances.par_iter_mut()) .enumerate() @@ -503,16 +513,22 @@ impl Restart { inertia = self.point_distances.iter().sum(); - // SAFETY: `sample.len() == m * d` with `m` labels, sums is - // `k * d` with `k` counts, and `d` is a multiple of 8 - // (guaranteed by Dimension). + let mut scratch = Scratch { + sums: &mut self.sums, + counts: &mut self.counts, + order: &mut self.order, + bounds: &mut self.bounds, + }; + + // SAFETY: `d` is a multiple of 8 (guaranteed by Dimension); + // every other requirement is checked by `accumulate_clusters` + // itself and panics rather than misbehaving. unsafe { accumulate_clusters( sample, &self.labels, Some(sample_inv_norms), - &mut self.sums, - &mut self.counts, + &mut scratch, d, ); } @@ -522,16 +538,17 @@ impl Restart { continue; } - let start = cluster * d; + let start = cluster * d.get(); // Normalization is scale-invariant, so the raw sum gives the same direction as the // average. - self.centroids[start..start + d].copy_from_slice(&self.sums[start..start + d]); + self.centroids[start..start + d.get()] + .copy_from_slice(&self.sums[start..start + d.get()]); // SAFETY: centroid rows have length `d`, and `d` is a multiple of 8 // (guaranteed by Dimension). unsafe { - kernel::normalize(&mut self.centroids[start..start + d]); + kernel::normalize(&mut self.centroids[start..start + d.get()]); } } @@ -586,15 +603,15 @@ impl Restart { } } - let point_start = farthest_idx * d; - let centroid_start = cluster * d; + let point_start = farthest_idx * d.get(); + let centroid_start = cluster * d.get(); - self.centroids[centroid_start..centroid_start + d] - .copy_from_slice(&sample[point_start..point_start + d]); + self.centroids[centroid_start..centroid_start + d.get()] + .copy_from_slice(&sample[point_start..point_start + d.get()]); // SAFETY: centroid rows have length `d`, a multiple of 8. unsafe { - kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d]); + kernel::normalize(&mut self.centroids[centroid_start..centroid_start + d.get()]); } self.labels[farthest_idx] = cluster as u16; @@ -605,53 +622,126 @@ impl Restart { } } +/// Borrowed scratch for [`accumulate_clusters`]. +struct Scratch<'ctx> { + /// Per-cluster accumulator, `k * d` elements. + sums: &'ctx mut [f32], + /// Per-cluster point count, `k` elements. + counts: &'ctx mut [usize], + /// Point indices grouped by cluster, one per labeled point. + order: &'ctx mut [usize], + /// Bucket cursors during the scatter, bucket boundaries after; + /// `k + 1` elements. + bounds: &'ctx mut [usize], +} + /// Recomputes per-cluster sums and counts from labeled points. /// -/// The result is impervious to any thread schedule order. +/// Points are grouped by cluster first (a stable counting sort on labels, +/// through `order` and `bounds`), then each cluster task walks only its own +/// bucket: one pass over the labels instead of one per cluster. Buckets +/// preserve ascending point order, so every cluster sum receives its +/// additions in a fixed order and the result is impervious to any thread +/// schedule order: the grouping runs sequentially on the calling thread, +/// and each sum is reduced by exactly one task over a fixed range. /// /// `inv_norms` supplies precomputed inverse norms; pass `None` to compute /// them on the fly. /// /// Zero-norm points are counted but contribute nothing to the sums. /// +/// `order` and `bounds` are grouping scratch; their contents on entry are +/// irrelevant. +/// +/// # Panics +/// +/// Panics if any label is not below `counts.len()`, or if the scratch and +/// input shapes are inconsistent: `order.len() != labels.len()`, +/// `bounds.len() != counts.len() + 1`, `sums.len() != counts.len() * d`, +/// or `inv_norms` (when provided) not one entry per label. +/// /// # Safety /// -/// * `points.len() == labels.len() * d` -/// * `sums.len() == counts.len() * d` -/// * `inv_norms`, when provided, has one entry per label -/// * `d` is a multiple of 8 +/// * `d` is a multiple of 8 (the SIMD kernels rely on it; every other requirement is checked at +/// runtime and panics) unsafe fn accumulate_clusters( points: &[f32], labels: &[u16], inv_norms: Option<&[f32]>, - sums: &mut [f32], - counts: &mut [usize], - d: usize, + Scratch { + sums, + counts, + order, + bounds, + }: &mut Scratch<'_>, + d: NonZero, ) { - // A `debug_assert` only: an `assert_unchecked` here would be sound (the - // length is a documented precondition), but to elide the - // `inv_norms[index]` bounds check the fact would have to survive into - // the rayon closure, and that check is noise next to the `d`-wide - // kernel call it precedes. - debug_assert!(inv_norms.is_none_or(|norms| norms.len() == labels.len())); + let d = d.get(); + + // We deliberately opt into checked indexing here. Profiling via performance counters on + // Apple M5 (which does have a large out-of-order execution window) showed that the fully + // checked version has negligible cost: an additional ~13 instructions and ~6 branches per + // point (+5.3% instructions), yet no increase in cycle count. This is because the check + // branches are >99.8% predicted, and the extra µops retire from issue slots that otherwise + // sit idle behind FMA and load latency (backend stall slots *drop* ~4%); the loop is + // memory/FMA-bound, not issue-bound. + // + // While the measurements are specific to Apple M5, the general trend should be transferable + // to other architectures. + // + // The checks buy panics-instead-of-UB for free, shrinking `# Safety` to the kernels' + // alignment requirement. Do not switch this back to unchecked indexing without new + // measurements. + assert!(inv_norms.is_none_or(|norms| norms.len() == labels.len())); + assert_eq!(order.len(), labels.len()); + assert_eq!(bounds.len(), counts.len() + 1); + assert!(!bounds.is_empty()); + assert_eq!(sums.len(), counts.len() * d); + + // 1. Histogram: the counts double as the bucket sizes, and the checked indexing doubles as + // validation: an out-of-range label panics here, before any scratch is written. + counts.fill(0); + for &label in labels { + counts[usize::from(label)] += 1; + } - // TODO(PERF): This isn't as performant as it could be, in theory we could first create a - // histogram, and then use that to dispatch over it. A potential optimization opportunity if - // ever required. + // 2. Bucket starts. During the scatter, `bounds[c + 1]` is cluster `c`'s write cursor; each + // cursor ends at its bucket end, leaving `bounds` as exactly the boundary array the gather + // needs: cluster `c` owns `order[bounds[c]..bounds[c + 1]]`. + bounds[0] = 0; + let mut running = 0; + for (bound, &count) in bounds[1..].iter_mut().zip(counts.iter()) { + *bound = running; + running += count; + } + + // 3. Stable scatter: visiting points in ascending index order keeps each bucket ascending, + // which pins the per-cluster addition order (and therefore the sums) regardless of how rayon + // schedules the gather. + for (index, &label) in labels.iter().enumerate() { + // Cluster `c`'s cursor starts at its bucket start and is bumped once + // per point labeled `c`, of which the histogram counted exactly + // `counts[c]`, so it stays below the bucket end (at most + // `order.len()`): for histogram-validated labels these accesses + // cannot panic. + let cursor = &mut bounds[usize::from(label) + 1]; + order[*cursor] = index; + *cursor += 1; + } + // Shared views for the parallel gather. + let order: &[usize] = order; + let bounds: &[usize] = bounds; + + // 4. Accumulate: one task per cluster, walking only its own bucket. sums.par_chunks_exact_mut(d) - .zip(counts.par_iter_mut()) - .enumerate() - .for_each(|(cluster, (sum, count))| { + .zip(bounds.par_array_windows::<2>()) + .for_each(|(sum, &[start, end])| { sum.fill(0.0); - *count = 0; - for (index, (point, &label)) in points.chunks_exact(d).zip(labels).enumerate() { - if usize::from(label) != cluster { - continue; - } - - *count += 1; + for &index in &order[start..end] { + let row = index * d; + let point = &points[row..row + d]; let inv_norm = inv_norms.map_or_else( || { @@ -686,18 +776,21 @@ unsafe fn accumulate_clusters( /// * `d` is a multiple of 8 unsafe fn label_chunk( centroids: &[f32], - k: NonZero, - d: usize, + k_nz: NonZero, + d_nz: NonZero, points: &[f32], labels: &mut [u16], ) { + let k = k_nz.get(); + let d = d_nz.get(); + let count = labels.len(); // SAFETY: each parallel chunk pairs `count` labels with `count * d` // floats of point data; `d` is a multiple of 8 (guaranteed by Dimension). unsafe { core::hint::assert_unchecked(points.len() == count * d); - core::hint::assert_unchecked(centroids.len() >= k.get() * d); + core::hint::assert_unchecked(centroids.len() >= k * d); core::hint::assert_unchecked(d.is_multiple_of(8)); } @@ -710,7 +803,7 @@ unsafe fn label_chunk( // SAFETY: each point length d, centroids length k*d, k > 0, // d a multiple of 8 (guaranteed by Dimension). - let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; + let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k_nz, d_nz) }; labels[i] = nearest[0].0; labels[i + 1] = nearest[1].0; @@ -722,7 +815,7 @@ unsafe fn label_chunk( while i < count { let point = &points[i * d..i * d + d]; // SAFETY: point length d, centroids length k*d, k > 0, d mult of 8. - let (label, _) = unsafe { nearest_centroid(point, 1.0, centroids, k, d) }; + let (label, _) = unsafe { nearest_centroid(point, 1.0, centroids, k_nz, d_nz) }; labels[i] = label; i += 1; } @@ -739,10 +832,12 @@ unsafe fn label_chunk( unsafe fn score_chunk( centroids: &[f32], k: NonZero, - d: usize, + d_nz: NonZero, points: &[f32], labels: &mut [u16], ) -> f32 { + let d = d_nz.get(); + debug_assert_eq!(points.len(), labels.len() * d); // SAFETY: The caller must ensure `points.len() == labels.len() * d`. @@ -762,7 +857,7 @@ unsafe fn score_chunk( // SAFETY: each point length d, centroids length k*d, k > 0, // d a multiple of 8 (guaranteed by Dimension). - let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d) }; + let nearest = unsafe { kernel::nearest4(p0, p1, p2, p3, centroids, k, d_nz) }; let ps = [p0, p1, p2, p3]; for offset in 0..4 { @@ -784,7 +879,7 @@ unsafe fn score_chunk( let inv_norm = if norm > 0.0 { norm.recip() } else { 0.0 }; // SAFETY: point length d, centroids length k*d, k > 0, d mult of 8. - let (label, distance) = unsafe { nearest_centroid(point, inv_norm, centroids, k, d) }; + let (label, distance) = unsafe { nearest_centroid(point, inv_norm, centroids, k, d_nz) }; labels[i] = label; inertia += distance; i += 1; @@ -806,7 +901,7 @@ unsafe fn reassign( centroids: &[f32], labels: &mut [u16], k: NonZero, - d: usize, + d: NonZero, chunk: usize, row_chunk: usize, ) { @@ -839,7 +934,7 @@ unsafe fn reassign_scored( centroids: &[f32], labels: &mut [u16], k: NonZero, - d: usize, + d: NonZero, chunk: usize, row_chunk: usize, ) -> f32 { @@ -861,15 +956,14 @@ unsafe fn reassign_scored( /// from the full population, and re-labels against the final centroids. /// Returns the full-data inertia. /// -/// `sums` and `counts` are accumulator scratch; their contents on entry are -/// irrelevant. +/// `scratch` contents on entry are irrelevant; its shape is validated by +/// [`accumulate_clusters`], which panics on any mismatch. /// /// # Safety /// /// * `x.len() == n * d` for some `n` /// * `clustering.centroids.len() == k * d` /// * `clustering.labels.len() == n` -/// * `sums.len() == k * d` and `counts.len() == k` /// * `d` is a multiple of 8 unsafe fn assign( x: &[f32], @@ -877,10 +971,9 @@ unsafe fn assign( k: NonZero, chunk: usize, row_chunk: usize, - sums: &mut [f32], - counts: &mut [usize], + scratch: &mut Scratch<'_>, ) -> f32 { - let d = clustering.dimension.get() as usize; + let d = NonZero::::from(clustering.dimension.value()); // 1. Label all points against the sample-fitted centroids. // SAFETY: forwarded from the caller. @@ -897,18 +990,19 @@ unsafe fn assign( } // 2. Recompute centroids from the full population. - // SAFETY: `x.len() == n * d` with `n` labels, sums is `k * d` with `k` - // counts, and `d` is a multiple of 8 (guaranteed by Dimension). + // SAFETY: `d` is a multiple of 8 (guaranteed by Dimension); every other + // requirement is checked by `accumulate_clusters` itself and panics + // rather than misbehaving. unsafe { - accumulate_clusters(x, &clustering.labels, None, sums, counts, d); + accumulate_clusters(x, &clustering.labels, None, scratch, d); } - for (cluster, count) in counts.iter_mut().enumerate() { + for (cluster, count) in scratch.counts.iter_mut().enumerate() { if *count == 0 { continue; } - let start = cluster * d; + let start = cluster * d.get(); #[expect( clippy::cast_possible_truncation, @@ -917,7 +1011,7 @@ unsafe fn assign( let centroid = clustering.centroid_mut(cluster as u16); // Normalization is scale-invariant, so the raw sum gives the same // direction as the average. - centroid.copy_from_slice(&sums[start..start + d]); + centroid.copy_from_slice(&scratch.sums[start..start + d.get()]); // SAFETY: centroid length d, a multiple of 8. unsafe { @@ -1028,7 +1122,7 @@ pub fn cluster(x: &[f32], dimension: Dimension, config: &Config) -> Clustering { .into_par_iter() .enumerate() .map(|(index, seed)| { - let mut restart = Restart::new(k, m, d); + let mut restart = Restart::new(k, m, dimension.value().into()); let inertia = restart.run(sample, chunk, row_chunk, &sample_inv_norms, seed, config); (inertia, index, restart) @@ -1038,31 +1132,35 @@ pub fn cluster(x: &[f32], dimension: Dimension, config: &Config) -> Clustering { // Reuse the winning restart's buffers: its centroids become the result // and its per-cluster accumulators serve the full-data recomputation, - // instead of allocating fresh ones. + // instead of allocating fresh ones. Only `order` needs a fresh, `n`-sized + // allocation (the restart's is sample-sized); this is the sole per-fit + // scratch allocation outside setup, and none happen per iteration. let Restart { centroids, mut sums, mut counts, + mut bounds, .. } = best.2; clustering.centroids = centroids; + let mut order = vec![0_usize; n]; + + let mut scratch = Scratch { + sums: &mut sums, + counts: &mut counts, + order: &mut order, + bounds: &mut bounds, + }; + // 3. assign points to clusters // SAFETY: `x.len() == n * d` (asserted above), `clustering.centroids.len() == k * d`, // `sums` and `counts` are the restart's `k * d` and `k` sized accumulators, - // and `d` is a multiple of 8 (guaranteed by Dimension). - clustering.inertia = unsafe { - assign( - x, - &mut clustering, - k, - chunk, - row_chunk, - &mut sums, - &mut counts, - ) - }; + // `order` was just allocated with `n` entries, `bounds` is the restart's + // `k + 1` boundary scratch, and `d` is a multiple of 8 (guaranteed by + // Dimension). + clustering.inertia = unsafe { assign(x, &mut clustering, k, chunk, row_chunk, &mut scratch) }; clustering } @@ -1565,7 +1663,7 @@ mod tests { // SAFETY: point has length D=64, centroids has length k*D, // k > 0, D is a multiple of 8. - let (got, _) = unsafe { nearest_centroid(&p, inv, ¢roids, k, D) }; + let (got, _) = unsafe { nearest_centroid(&p, inv, ¢roids, k, nz!(D)) }; assert_eq!( got, brute_nearest_cosine(&p, ¢roids, k), @@ -1587,9 +1685,9 @@ mod tests { // SAFETY: point has length D=64, centroids has length k*D, // k > 0, D is a multiple of 8. - let (a, _) = unsafe { nearest_centroid(&p, 1.0, ¢roids, k, D) }; + let (a, _) = unsafe { nearest_centroid(&p, 1.0, ¢roids, k, nz!(D)) }; // SAFETY: same preconditions. - let (b, _) = unsafe { nearest_centroid(&p, 0.123, ¢roids, k, D) }; + let (b, _) = unsafe { nearest_centroid(&p, 0.123, ¢roids, k, nz!(D)) }; assert_eq!(a, b, "inv_norm must not change the selected centroid"); } } @@ -1614,4 +1712,139 @@ mod tests { assert!(result.centroids.iter().all(|v| v.is_finite())); assert!(result.labels.iter().all(|&l| l < 5)); } + + /// Sequential reference for [`accumulate_clusters`]: one pass over the + /// points in ascending index order, adding each to its cluster's sum + /// with the same kernels. Bit-identical to the parallel version by + /// construction if (and only if) the grouping preserves per-cluster + /// ascending order, which is exactly the property under test. + fn accumulate_reference( + points: &[f32], + labels: &[u16], + inv_norms: Option<&[f32]>, + k: usize, + d: usize, + ) -> (Vec, Vec) { + let mut sums = vec![0.0_f32; k * d]; + let mut counts = vec![0_usize; k]; + + for (index, (point, &label)) in points.chunks_exact(d).zip(labels).enumerate() { + let cluster = usize::from(label); + counts[cluster] += 1; + + let inv_norm = inv_norms.map_or_else( + || { + // SAFETY: `point` has length `d`, a multiple of 8. + let norm = unsafe { kernel::dot(point, point) }.sqrt(); + + if norm > 0.0 { norm.recip() } else { 0.0 } + }, + |inv_norms| inv_norms[index], + ); + + if inv_norm == 0.0 { + continue; + } + + // SAFETY: sum rows and `point` have length `d`, a multiple of 8. + unsafe { + kernel::add_scaled_into(&mut sums[cluster * d..(cluster + 1) * d], point, inv_norm); + } + } + + (sums, counts) + } + + #[test] + fn accumulate_clusters_matches_sequential_reference_bitwise() { + // Half the points land in cluster 0 (skew), the rest spread over + // 1..=5; cluster 6 stays empty. + const PATTERN: [u16; 10] = [0, 1, 0, 2, 0, 3, 0, 4, 0, 5]; + let n = 203; + let d = 32; + let k = 7_usize; + + let mut rng = Xoshiro256PlusPlus::seed_from_u64(31); + let mut points = vec![0.0_f32; n * d]; + for (i, row) in points.chunks_exact_mut(d).enumerate() { + if i % 17 == 0 { + continue; // leave all-zero: counted, but contributes nothing + } + for v in row.iter_mut() { + *v = rng.random_range(-1.0..1.0); + } + } + + let labels: Vec = (0..n).map(|i| PATTERN[i % PATTERN.len()]).collect(); + + let inv_norms: Vec = points + .chunks_exact(d) + .map(|row| { + let norm = l2(row); + if norm > 0.0 { norm.recip() } else { 0.0 } + }) + .collect(); + + for inv in [None, Some(inv_norms.as_slice())] { + // Sentinels: the accumulator must overwrite all of its outputs. + let mut sums = vec![f32::NAN; k * d]; + let mut counts = vec![usize::MAX; k]; + let mut order = vec![usize::MAX; n]; + let mut bounds = vec![usize::MAX; k + 1]; + + let mut scratch = Scratch { + sums: &mut sums, + counts: &mut counts, + order: &mut order, + bounds: &mut bounds, + }; + + // SAFETY: `points` is `n * d` floats with `n` labels, all labels + // are drawn from `PATTERN` and below `k == 7`, sums are `k * d` + // with `k` counts, inv norms (when provided) have one entry per + // label, order has `n` entries, bounds has `k + 1`, and + // `d == 32` is a multiple of 8. + unsafe { + accumulate_clusters(&points, &labels, inv, &mut scratch, nz!(32)); + } + + let (expected_sums, expected_counts) = + accumulate_reference(&points, &labels, inv, k, d); + + assert_eq!(counts, expected_counts); + assert_eq!(counts[6], 0, "cluster 6 must stay empty"); + assert_eq!(counts.iter().sum::(), n, "every point counted"); + + for (i, (sum, expected)) in sums.iter().zip(&expected_sums).enumerate() { + assert_eq!( + sum.to_bits(), + expected.to_bits(), + "sums diverge at element {i}: {sum} vs {expected}" + ); + } + } + } + + /// The determinism contract: identical bits regardless of pool width. + /// Guards the accumulate/assignment structure against rewrites whose + /// float reduction order depends on rayon's scheduling. + #[test] + fn cluster_bitwise_identical_across_pool_sizes() { + let (data, _) = make_blobs::<8>(40, 4, 2024); + let config = Config::for_k_with_seed(4, 7); + + let single = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .expect("single-thread pool should build") + .install(|| cluster(&data, dim(8), &config)); + let multi = rayon::ThreadPoolBuilder::new() + .num_threads(8) + .build() + .expect("multi-thread pool should build") + .install(|| cluster(&data, dim(8), &config)); + + assert_eq!(single.labels, multi.labels); + assert_eq!(single.centroids, multi.centroids); + } } diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/store/src/embedding/kernel.rs index 67f99a66831..f1881eb0ab9 100644 --- a/libs/@local/graph/store/src/embedding/kernel.rs +++ b/libs/@local/graph/store/src/embedding/kernel.rs @@ -345,8 +345,11 @@ pub unsafe fn nearest4( p3: &[f32], centroids: &[f32], k: NonZero, - d: usize, + d: NonZero, ) -> [(u16, f32); 4] { + let d = d.get(); + let k = k.get(); + let mut best_dot = [f32::NEG_INFINITY; 4]; let mut best_idx = [0_u16; 4]; @@ -356,12 +359,12 @@ pub unsafe fn nearest4( core::hint::assert_unchecked(p0.len() == p1.len()); core::hint::assert_unchecked(p0.len() == p2.len()); core::hint::assert_unchecked(p0.len() == p3.len()); - core::hint::assert_unchecked(centroids.len() >= k.get() * d); + core::hint::assert_unchecked(centroids.len() >= k * d); core::hint::assert_unchecked(d.is_multiple_of(8)); } let mut j = 0; - while j + 2 <= k.get() { + while j + 2 <= k { // SAFETY: `j + 2 <= k` and `centroids.len() >= k * d`, so both // slices `[j*d .. (j+2)*d]` are in-bounds. let c0 = unsafe { centroids.get_unchecked(j * d..j * d + d) }; @@ -389,7 +392,7 @@ pub unsafe fn nearest4( } // Handle odd k: one remaining centroid via the 4x1 tile. - if j < k.get() { + if j < k { let c = ¢roids[j * d..j * d + d]; // SAFETY: all five slices have length `d`, a multiple of 8. let dots = unsafe { micro_4x1(p0, p1, p2, p3, c) }; @@ -719,7 +722,13 @@ mod tests { // all point slices have length d. let got = unsafe { nearest4( - &points[0], &points[1], &points[2], &points[3], ¢roids, k, d, + &points[0], + &points[1], + &points[2], + &points[3], + ¢roids, + k, + nz!(8), ) }; @@ -749,7 +758,13 @@ mod tests { // all point slices have length d. let got = unsafe { nearest4( - &points[0], &points[1], &points[2], &points[3], ¢roids, k, d, + &points[0], + &points[1], + &points[2], + &points[3], + ¢roids, + k, + nz!(8), ) }; @@ -770,7 +785,7 @@ mod tests { // SAFETY: d=8 (multiple of 8), k=1 > 0, centroids has length d, // all point slices have length d. - let got = unsafe { nearest4(&p0, &p1, &p2, &p3, ¢roids, nz!(1), d) }; + let got = unsafe { nearest4(&p0, &p1, &p2, &p3, ¢roids, nz!(1), nz!(8)) }; assert_eq!(got[0].0, 0); assert_eq!(got[1].0, 0); From 0138056066f9be3d7e6e436e816cb37c4df5fb9b Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:31:25 +0200 Subject: [PATCH 24/38] chore: docs --- libs/@local/graph/store/src/embedding/clustering.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/store/src/embedding/clustering.rs index 3b92fb711cf..8b3d011a483 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/store/src/embedding/clustering.rs @@ -662,8 +662,7 @@ struct Scratch<'ctx> { /// /// # Safety /// -/// * `d` is a multiple of 8 (the SIMD kernels rely on it; every other requirement is checked at -/// runtime and panics) +/// * `d` is a multiple of 8 unsafe fn accumulate_clusters( points: &[f32], labels: &[u16], From 2dafd88ca1ea824d8c8afeb70fa4d67c134cae46 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:23:26 +0200 Subject: [PATCH 25/38] feat: move into embeddings crate --- Cargo.lock | 11 +- libs/@local/graph/embeddings/Cargo.toml | 12 ++ .../benches/clustering.rs} | 125 +++++++++--------- .../src}/clustering.rs | 9 ++ .../embedding => embeddings/src}/dimension.rs | 20 ++- .../embedding => embeddings/src}/kernel.rs | 9 ++ libs/@local/graph/embeddings/src/lib.rs | 23 +++- libs/@local/graph/postgres-store/Cargo.toml | 37 +++--- .../store/postgres/knowledge/entity/mod.rs | 6 +- libs/@local/graph/store/Cargo.toml | 15 +-- libs/@local/graph/store/src/embedding/mod.rs | 16 --- libs/@local/graph/store/src/lib.rs | 7 +- 12 files changed, 157 insertions(+), 133 deletions(-) rename libs/@local/graph/{store/benches/embedding.rs => embeddings/benches/clustering.rs} (66%) rename libs/@local/graph/{store/src/embedding => embeddings/src}/clustering.rs (99%) rename libs/@local/graph/{store/src/embedding => embeddings/src}/dimension.rs (79%) rename libs/@local/graph/{store/src/embedding => embeddings/src}/kernel.rs (99%) delete mode 100644 libs/@local/graph/store/src/embedding/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 3045f2d5588..c4912b58916 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3678,9 +3678,14 @@ dependencies = [ name = "hash-graph-embeddings" version = "0.0.0" dependencies = [ + "codspeed-criterion-compat", + "darwin-kperf-criterion", "derive_more", "error-stack", "hash-graph-types", + "rand 0.10.1", + "rand_xoshiro", + "rayon", "reqwest", "reqwest-middleware", "reqwest-retry", @@ -3771,6 +3776,7 @@ dependencies = [ "futures-sink", "hash-codec", "hash-graph-authorization", + "hash-graph-embeddings", "hash-graph-migrations", "hash-graph-store", "hash-graph-temporal-versioning", @@ -3804,8 +3810,6 @@ name = "hash-graph-store" version = "0.0.0" dependencies = [ "bytes", - "codspeed-criterion-compat", - "darwin-kperf-criterion", "derive-where", "derive_more", "error-stack", @@ -3818,9 +3822,6 @@ dependencies = [ "hash-temporal-client", "insta", "postgres-types", - "rand 0.10.1", - "rand_xoshiro", - "rayon", "serde", "serde_json", "simple-mermaid", diff --git a/libs/@local/graph/embeddings/Cargo.toml b/libs/@local/graph/embeddings/Cargo.toml index 907e2eb5420..d903507bb96 100644 --- a/libs/@local/graph/embeddings/Cargo.toml +++ b/libs/@local/graph/embeddings/Cargo.toml @@ -17,6 +17,9 @@ error-stack = { workspace = true } # Private third-party dependencies derive_more = { workspace = true, features = ["display", "error"] } +rand = { workspace = true } +rand_xoshiro = { workspace = true } +rayon = { workspace = true } reqwest = { workspace = true } reqwest-middleware = { workspace = true, features = ["json"] } reqwest-retry = { workspace = true } @@ -25,5 +28,14 @@ serde = { workspace = true, features = ["derive"] } simple-mermaid = { workspace = true } tracing = { workspace = true } +[dev-dependencies] +codspeed-criterion-compat = { workspace = true } +darwin-kperf-criterion = { workspace = true, features = ["codspeed"] } + +[[bench]] +name = "clustering" +harness = false + + [lints] workspace = true diff --git a/libs/@local/graph/store/benches/embedding.rs b/libs/@local/graph/embeddings/benches/clustering.rs similarity index 66% rename from libs/@local/graph/store/benches/embedding.rs rename to libs/@local/graph/embeddings/benches/clustering.rs index d5dc2268cc3..24658df6723 100644 --- a/libs/@local/graph/store/benches/embedding.rs +++ b/libs/@local/graph/embeddings/benches/clustering.rs @@ -23,33 +23,39 @@ drop-tightening warning originates inside `criterion_group!`" )] -use core::hint::black_box; +use core::{hint::black_box, num::NonZero}; use codspeed_criterion_compat::{ BenchmarkId, Criterion, criterion_group, criterion_main, measurement::Measurement, }; -use hash_graph_store::embedding::{ +use hash_graph_embeddings::{ + D256, D1536, D3072, Dimension, clustering::{Config, cluster}, - dimension::Dimension, kernel, }; -use rand::{RngExt as _, SeedableRng as _}; +use rand::{RngExt as _, SeedableRng as _, distr::Uniform}; use rand_xoshiro::Xoshiro256PlusPlus; +macro_rules! nz { + ($expr:expr) => { + const { ::core::num::NonZero::new($expr).unwrap() } + }; +} + /// Uniform random values in `[-1, 1)`. -fn random_vec(len: usize, seed: u64) -> Vec { - let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); - core::iter::repeat_with(|| rng.random_range(-1.0..1.0)) - .take(len) +fn random_vec(n: usize, seed: u64) -> Vec { + let rng = Xoshiro256PlusPlus::seed_from_u64(seed); + rng.sample_iter(Uniform::new(-1.0, 1.0).expect("uniform range is non-empty")) + .take(n) .collect() } /// Uniform random values in `[0.1, 1)`, guaranteed positive so repeated /// accumulation saturates at infinity instead of producing NaNs. -fn random_positive_vec(len: usize, seed: u64) -> Vec { - let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed); - core::iter::repeat_with(|| rng.random_range(0.1..1.0)) - .take(len) +fn random_positive_vec(n: usize, seed: u64) -> Vec { + let rng = Xoshiro256PlusPlus::seed_from_u64(seed); + rng.sample_iter(Uniform::new(0.1, 1.0).expect("uniform range is non-empty")) + .take(n) .collect() } @@ -72,16 +78,16 @@ fn blobs(points_per_cluster: usize, k: usize, d: usize, seed: u64) -> Vec { data } -const KERNEL_DIMS: &[usize] = &[256, 1536, 3072]; +const KERNEL_DIMS: [Dimension; 3] = [D256, D1536, D3072]; fn bench_dot(criterion: &mut Criterion) { - let mut group = criterion.benchmark_group("embedding/kernel/dot"); + let mut group = criterion.benchmark_group("kernel/dot"); - for &d in KERNEL_DIMS { - let lhs = random_vec(d, 1); - let rhs = random_vec(d, 2); + for dim in KERNEL_DIMS { + let lhs = random_vec(dim.get() as usize, 1); + let rhs = random_vec(dim.get() as usize, 2); - group.bench_with_input(BenchmarkId::from_parameter(d), &d, |bencher, _| { + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { // SAFETY: both slices have length `d`, a multiple of 8. bencher.iter(|| unsafe { kernel::dot(black_box(&lhs), black_box(&rhs)) }); }); @@ -91,13 +97,13 @@ fn bench_dot(criterion: &mut Criterion) { } fn bench_add_scaled_into(criterion: &mut Criterion) { - let mut group = criterion.benchmark_group("embedding/kernel/add_scaled_into"); + let mut group = criterion.benchmark_group("kernel/add_scaled_into"); - for &d in KERNEL_DIMS { - let src = random_positive_vec(d, 3); - let mut dst = random_positive_vec(d, 4); + for dim in KERNEL_DIMS { + let src = random_positive_vec(dim.get() as usize, 3); + let mut dst = random_positive_vec(dim.get() as usize, 4); - group.bench_with_input(BenchmarkId::from_parameter(d), &d, |bencher, _| { + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { // SAFETY: both slices have length `d`, a multiple of 8. bencher.iter(|| unsafe { kernel::add_scaled_into(black_box(&mut dst), black_box(&src), black_box(0.5)); @@ -109,21 +115,22 @@ fn bench_add_scaled_into(criterion: &mut Criterion) { } fn bench_micro_4x2(criterion: &mut Criterion) { - let mut group = criterion.benchmark_group("embedding/kernel/micro_4x2"); + let mut group = criterion.benchmark_group("kernel/micro_4x2"); - for &d in KERNEL_DIMS { - let points: Vec> = (0..4).map(|seed| random_vec(d, 10 + seed)).collect(); - let c0 = random_vec(d, 20); - let c1 = random_vec(d, 21); + for dim in KERNEL_DIMS { + let [p0, p1, p2, p3] = + core::array::from_fn(|index| random_vec(dim.get() as usize, 10 + index as u64)); + let c0 = random_vec(dim.get() as usize, 20); + let c1 = random_vec(dim.get() as usize, 21); - group.bench_with_input(BenchmarkId::from_parameter(d), &d, |bencher, _| { + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { // SAFETY: all six slices have length `d`, a multiple of 8. bencher.iter(|| unsafe { kernel::micro_4x2( - black_box(&points[0]), - black_box(&points[1]), - black_box(&points[2]), - black_box(&points[3]), + black_box(&p0), + black_box(&p1), + black_box(&p2), + black_box(&p3), black_box(&c0), black_box(&c1), ) @@ -134,41 +141,36 @@ fn bench_micro_4x2(criterion: &mut Criterion) { group.finish(); } -macro_rules! nz { - ($expr:expr) => { - const { ::core::num::NonZero::new($expr).unwrap() } - }; -} - fn bench_nearest4(criterion: &mut Criterion) { - let mut group = criterion.benchmark_group("embedding/kernel/nearest4"); + let mut group = criterion.benchmark_group("kernel/nearest4"); // k = 15 exercises the odd-k remainder path. - for &(d, k) in &[ - (nz!(256), nz!(15)), - (nz!(256), nz!(16)), - (nz!(256), nz!(64)), - (nz!(1536), nz!(16)), - (nz!(3072), nz!(16)), + for &(dim, k) in &[ + (D256, nz!(15)), + (D256, nz!(16)), + (D256, nz!(64)), + (D1536, nz!(16)), + (D3072, nz!(16)), ] { - let points: Vec> = (0..4).map(|seed| random_vec(d.get(), 30 + seed)).collect(); - let centroids = random_vec(k.get() * d.get(), 40); + let [p0, p1, p2, p3] = + core::array::from_fn(|index| random_vec(dim.get() as usize, 30 + index as u64)); + let centroids = random_vec(k.get() * dim.get() as usize, 40); group.bench_with_input( - BenchmarkId::new(format!("d{d}"), k), - &(d, k), + BenchmarkId::new(format!("d{dim}"), k), + &(dim, k), |bencher, _| { // SAFETY: point slices have length `d` (multiple of 8), // centroids has length `k * d`, and `k > 0`. bencher.iter(|| unsafe { kernel::nearest4( - black_box(&points[0]), - black_box(&points[1]), - black_box(&points[2]), - black_box(&points[3]), + black_box(&p0), + black_box(&p1), + black_box(&p2), + black_box(&p3), black_box(¢roids), black_box(k), - black_box(d), + black_box(NonZero::from(dim.value())), ) }); }, @@ -179,15 +181,12 @@ fn bench_nearest4(criterion: &mut Criterion) { } fn bench_cluster(criterion: &mut Criterion) { - rayon::ThreadPoolBuilder::new() + let pool = rayon::ThreadPoolBuilder::new() .num_threads(1) - .build_global() + .build() .expect("should be built exactly once"); - let mut group = criterion.benchmark_group("embedding/cluster"); - group.sample_size(10); - - let dimension = Dimension::new(256).expect("256 is a positive multiple of 8"); + let mut group = criterion.benchmark_group("cluster"); // (n, k): n = 10k exercises the subsampled fit (m = 8192) plus the // full-data refinement; n = 50k shifts the weight onto the full-data @@ -205,7 +204,9 @@ fn bench_cluster(criterion: &mut Criterion) { BenchmarkId::new(format!("n{n}_d256"), k), &(n, k), |bencher, _| { - bencher.iter(|| cluster(black_box(&data), black_box(dimension), &config)); + pool.install(|| { + bencher.iter(|| cluster(black_box(&data), black_box(D256), &config)); + }); }, ); } diff --git a/libs/@local/graph/store/src/embedding/clustering.rs b/libs/@local/graph/embeddings/src/clustering.rs similarity index 99% rename from libs/@local/graph/store/src/embedding/clustering.rs rename to libs/@local/graph/embeddings/src/clustering.rs index 8b3d011a483..45111c4b8f6 100644 --- a/libs/@local/graph/store/src/embedding/clustering.rs +++ b/libs/@local/graph/embeddings/src/clustering.rs @@ -1,3 +1,12 @@ +#![expect( + unsafe_code, + clippy::indexing_slicing, + clippy::float_arithmetic, + clippy::min_ident_chars, + clippy::many_single_char_names, + reason = "Single-char idents (k, n, m, d, x) are standard mathematical notation for \ + clustering." +)] use alloc::borrow::Cow; use core::{cmp, num::NonZero}; use std::collections::HashSet; diff --git a/libs/@local/graph/store/src/embedding/dimension.rs b/libs/@local/graph/embeddings/src/dimension.rs similarity index 79% rename from libs/@local/graph/store/src/embedding/dimension.rs rename to libs/@local/graph/embeddings/src/dimension.rs index 0ba85516ee3..3701ee85a91 100644 --- a/libs/@local/graph/store/src/embedding/dimension.rs +++ b/libs/@local/graph/embeddings/src/dimension.rs @@ -1,4 +1,4 @@ -use core::num::NonZero; +use core::{fmt, fmt::Display, num::NonZero}; /// An embedding vector dimension, guaranteed to be a positive multiple of 8. /// @@ -39,6 +39,12 @@ impl Dimension { } } +impl Display for Dimension { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(&self.get(), f) + } +} + pub const D128: Dimension = Dimension(NonZero::new(128).unwrap()); pub const D256: Dimension = Dimension(NonZero::new(256).unwrap()); pub const D512: Dimension = Dimension(NonZero::new(512).unwrap()); @@ -51,10 +57,10 @@ mod tests { #[test] fn valid_multiples_of_8() { - for v in [8, 16, 24, 128, 256, 3072] { + for value in [8, 16, 24, 128, 256, 3072] { assert!( - Dimension::new(v).is_some(), - "{v} should be a valid dimension" + Dimension::new(value).is_some(), + "{value} should be a valid dimension" ); } } @@ -66,10 +72,10 @@ mod tests { #[test] fn non_multiples_of_8_rejected() { - for v in [1, 2, 3, 4, 5, 6, 7, 9, 10, 15, 17, 100, 3071] { + for value in [1, 2, 3, 4, 5, 6, 7, 9, 10, 15, 17, 100, 3071] { assert!( - Dimension::new(v).is_none(), - "{v} should not be a valid dimension" + Dimension::new(value).is_none(), + "{value} should not be a valid dimension" ); } } diff --git a/libs/@local/graph/store/src/embedding/kernel.rs b/libs/@local/graph/embeddings/src/kernel.rs similarity index 99% rename from libs/@local/graph/store/src/embedding/kernel.rs rename to libs/@local/graph/embeddings/src/kernel.rs index f1881eb0ab9..13040e483d6 100644 --- a/libs/@local/graph/store/src/embedding/kernel.rs +++ b/libs/@local/graph/embeddings/src/kernel.rs @@ -1,8 +1,17 @@ +#![expect( + unsafe_code, + clippy::indexing_slicing, + clippy::float_arithmetic, + clippy::min_ident_chars, + reason = "Single-char idents (k, n, m, d, x) are standard mathematical notation for \ + clustering." +)] #![expect( clippy::inline_always, reason = "while usually discouraged, SIMD operations need to be inlined, as otherwise we \ spill SIMD registers, see the SIMD documentation." )] + use core::{ num::NonZero, simd::{Simd, f32x8, num::SimdFloat as _}, diff --git a/libs/@local/graph/embeddings/src/lib.rs b/libs/@local/graph/embeddings/src/lib.rs index a5e819a016f..0351f0a0f81 100644 --- a/libs/@local/graph/embeddings/src/lib.rs +++ b/libs/@local/graph/embeddings/src/lib.rs @@ -4,18 +4,33 @@ //! //! ## Workspace dependencies #![cfg_attr(doc, doc = simple_mermaid::mermaid!("../docs/dependency-diagram.mmd"))] +#![feature( + // Library Features + portable_simd, + integer_widen_truncate +)] -pub use self::{ - error::EmbeddingError, - openai::{OpenAiEmbeddingClient, OpenAiEmbeddingClientConfig}, -}; +extern crate alloc; mod error; mod openai; +pub mod clustering; +mod dimension; +// Hidden from docs: the kernel is an implementation detail, exposed only so +// the `embedding` bench target can measure it in isolation. +#[doc(hidden)] +pub mod kernel; + use error_stack::Report; use hash_graph_types::Embedding; +pub use self::{ + dimension::{D128, D256, D512, D1536, D3072, Dimension}, + error::EmbeddingError, + openai::{OpenAiEmbeddingClient, OpenAiEmbeddingClientConfig}, +}; + /// Generates embedding vectors for text inputs. /// /// Implementations call out to an embedding provider (e.g. OpenAI). The generated embeddings are diff --git a/libs/@local/graph/postgres-store/Cargo.toml b/libs/@local/graph/postgres-store/Cargo.toml index a96ddd34dc6..c95229feb08 100644 --- a/libs/@local/graph/postgres-store/Cargo.toml +++ b/libs/@local/graph/postgres-store/Cargo.toml @@ -31,24 +31,25 @@ hash-status = { workspace = true } type-system = { workspace = true, features = ["postgres"] } # Private third-party dependencies -async-scoped = { workspace = true, features = ["use-tokio"] } -bytes = { workspace = true } -clap = { workspace = true, optional = true, features = ["derive", "env"] } -derive-where = { workspace = true } -derive_more = { workspace = true } -dotenv-flow = { workspace = true } -futures = { workspace = true } -postgres-types = { workspace = true, features = ["derive", "with-serde_json-1"] } -refinery = { workspace = true, features = ["tokio-postgres"] } -regex = { workspace = true } -semver = { workspace = true, features = ["serde"] } -serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } -simple-mermaid = { workspace = true } -time = { workspace = true } -tracing = { workspace = true } -utoipa = { workspace = true, optional = true, features = ["uuid"] } -uuid = { workspace = true, features = ["v4", "serde"] } +async-scoped = { workspace = true, features = ["use-tokio"] } +bytes = { workspace = true } +clap = { workspace = true, optional = true, features = ["derive", "env"] } +derive-where = { workspace = true } +derive_more = { workspace = true } +dotenv-flow = { workspace = true } +futures = { workspace = true } +hash-graph-embeddings.workspace = true +postgres-types = { workspace = true, features = ["derive", "with-serde_json-1"] } +refinery = { workspace = true, features = ["tokio-postgres"] } +regex = { workspace = true } +semver = { workspace = true, features = ["serde"] } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +simple-mermaid = { workspace = true } +time = { workspace = true } +tracing = { workspace = true } +utoipa = { workspace = true, optional = true, features = ["uuid"] } +uuid = { workspace = true, features = ["v4", "serde"] } [dev-dependencies] hash-graph-migrations = { workspace = true } diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index d69d276cfaf..719c23e8d00 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -17,8 +17,8 @@ use hash_graph_authorization::policies::{ resource::{EntityResourceConstraint, ResourceConstraint}, store::{PolicyCreationParams, PrincipalStore as _}, }; +use hash_graph_embeddings::Dimension; use hash_graph_store::{ - embedding::dimension::Dimension, entity::{ ClusterEntitiesParams, ClusterEntitiesResponse, CreateEntityParams, DeleteEntitiesParams, DeletionSummary, EmptyEntityTypes, EntityCluster, EntityPermissions, EntityQueryCursor, @@ -2740,7 +2740,7 @@ where }); } - let config = hash_graph_store::embedding::clustering::Config::for_k_with_seed( + let config = hash_graph_embeddings::clustering::Config::for_k_with_seed( params.cluster_count, params.seed.unwrap_or_else(|| { std::time::SystemTime::UNIX_EPOCH @@ -2757,7 +2757,7 @@ where ); let result = tokio::task::spawn_blocking(move || { - hash_graph_store::embedding::clustering::cluster(&flat, dimension, &config) + hash_graph_embeddings::clustering::cluster(&flat, dimension, &config) }) .await .change_context(ClusterError::Store)?; diff --git a/libs/@local/graph/store/Cargo.toml b/libs/@local/graph/store/Cargo.toml index afd99bc632c..825638c1e35 100644 --- a/libs/@local/graph/store/Cargo.toml +++ b/libs/@local/graph/store/Cargo.toml @@ -29,9 +29,6 @@ bytes = { workspace = true, optional = true } derive-where = { workspace = true } derive_more = { workspace = true, features = ["display", "error"] } futures = { workspace = true } -rand = { workspace = true } -rand_xoshiro = { workspace = true } -rayon = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } simple-mermaid = { workspace = true } @@ -40,20 +37,14 @@ tracing = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] -codspeed-criterion-compat = { workspace = true } -darwin-kperf-criterion = { workspace = true, features = ["codspeed"] } -hash-codegen = { workspace = true } -insta = { workspace = true } -tokio = { workspace = true, features = ["macros"] } +hash-codegen = { workspace = true } +insta = { workspace = true } +tokio = { workspace = true, features = ["macros"] } [[test]] name = "codegen" required-features = ["codegen"] -[[bench]] -name = "embedding" -harness = false - [features] codegen = ["dep:specta", "type-system/codegen", "hash-graph-authorization/codegen"] utoipa = ["hash-graph-temporal-versioning/utoipa", "type-system/utoipa", "dep:utoipa"] diff --git a/libs/@local/graph/store/src/embedding/mod.rs b/libs/@local/graph/store/src/embedding/mod.rs deleted file mode 100644 index d184500214f..00000000000 --- a/libs/@local/graph/store/src/embedding/mod.rs +++ /dev/null @@ -1,16 +0,0 @@ -#![expect( - unsafe_code, - clippy::indexing_slicing, - clippy::float_arithmetic, - clippy::min_ident_chars, - clippy::many_single_char_names, - reason = "Single-char idents (k, n, m, d, x) are standard mathematical notation for \ - clustering." -)] - -pub mod clustering; -pub mod dimension; -// Hidden from docs: the kernel is an implementation detail, exposed only so -// the `embedding` bench target can measure it in isolation. -#[doc(hidden)] -pub mod kernel; diff --git a/libs/@local/graph/store/src/lib.rs b/libs/@local/graph/store/src/lib.rs index 668c46b20b8..15a1b96c52f 100644 --- a/libs/@local/graph/store/src/lib.rs +++ b/libs/@local/graph/store/src/lib.rs @@ -4,11 +4,7 @@ #![cfg_attr(doc, doc = simple_mermaid::mermaid!("../docs/dependency-diagram.mmd"))] #![feature( // Language Features - impl_trait_in_assoc_type, - - // Library Features, - portable_simd, - integer_widen_truncate, + impl_trait_in_assoc_type )] #![cfg_attr(test, feature( // Language Features @@ -27,7 +23,6 @@ pub mod oauth_provider; pub mod property_type; pub mod user_deletion; -pub mod embedding; pub mod error; pub mod filter; pub mod migration; From aff33173f97be4813a58d4d393f67b7976bfb9c0 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:28:19 +0200 Subject: [PATCH 26/38] fix: CI --- libs/@local/graph/embeddings/package.json | 6 +++- libs/@local/graph/postgres-store/Cargo.toml | 38 ++++++++++----------- libs/@local/graph/store/package.json | 3 -- 3 files changed, 24 insertions(+), 23 deletions(-) diff --git a/libs/@local/graph/embeddings/package.json b/libs/@local/graph/embeddings/package.json index 1ed4d31f77d..15d7db0be60 100644 --- a/libs/@local/graph/embeddings/package.json +++ b/libs/@local/graph/embeddings/package.json @@ -4,9 +4,13 @@ "private": true, "license": "AGPL-3", "scripts": { + "build:codspeed": "cargo codspeed build -p hash-graph-embeddings", "doc:dependency-diagram": "cargo run -p hash-repo-chores -- dependency-diagram --output docs/dependency-diagram.mmd --root hash-graph-embeddings --root-deps-and-dependents --link-mode non-roots --include-dev-deps --include-build-deps --logging-console-level info", "fix:clippy": "just clippy --fix", - "lint:clippy": "just clippy" + "lint:clippy": "just clippy", + "test:codspeed": "cargo codspeed run -p hash-graph-embeddings", + "test:miri": "cargo miri nextest run -- kernel clustering::tests::nearest_centroid_argmax_independent_of_inv_norm", + "test:unit": "mise run test:unit @rust/hash-graph-embeddings" }, "dependencies": { "@rust/error-stack": "workspace:*", diff --git a/libs/@local/graph/postgres-store/Cargo.toml b/libs/@local/graph/postgres-store/Cargo.toml index c95229feb08..114c71642d4 100644 --- a/libs/@local/graph/postgres-store/Cargo.toml +++ b/libs/@local/graph/postgres-store/Cargo.toml @@ -25,31 +25,31 @@ tokio-postgres = { workspace = true, public = true } # Private workspace dependencies error-stack = { workspace = true, features = ["std", "serde", "unstable"] } hash-codec = { workspace = true, features = ["numeric", "postgres"] } +hash-graph-embeddings = { workspace = true } hash-graph-temporal-versioning = { workspace = true, features = ["postgres"] } hash-graph-types = { workspace = true, features = ["postgres"] } hash-status = { workspace = true } type-system = { workspace = true, features = ["postgres"] } # Private third-party dependencies -async-scoped = { workspace = true, features = ["use-tokio"] } -bytes = { workspace = true } -clap = { workspace = true, optional = true, features = ["derive", "env"] } -derive-where = { workspace = true } -derive_more = { workspace = true } -dotenv-flow = { workspace = true } -futures = { workspace = true } -hash-graph-embeddings.workspace = true -postgres-types = { workspace = true, features = ["derive", "with-serde_json-1"] } -refinery = { workspace = true, features = ["tokio-postgres"] } -regex = { workspace = true } -semver = { workspace = true, features = ["serde"] } -serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } -simple-mermaid = { workspace = true } -time = { workspace = true } -tracing = { workspace = true } -utoipa = { workspace = true, optional = true, features = ["uuid"] } -uuid = { workspace = true, features = ["v4", "serde"] } +async-scoped = { workspace = true, features = ["use-tokio"] } +bytes = { workspace = true } +clap = { workspace = true, optional = true, features = ["derive", "env"] } +derive-where = { workspace = true } +derive_more = { workspace = true } +dotenv-flow = { workspace = true } +futures = { workspace = true } +postgres-types = { workspace = true, features = ["derive", "with-serde_json-1"] } +refinery = { workspace = true, features = ["tokio-postgres"] } +regex = { workspace = true } +semver = { workspace = true, features = ["serde"] } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +simple-mermaid = { workspace = true } +time = { workspace = true } +tracing = { workspace = true } +utoipa = { workspace = true, optional = true, features = ["uuid"] } +uuid = { workspace = true, features = ["v4", "serde"] } [dev-dependencies] hash-graph-migrations = { workspace = true } diff --git a/libs/@local/graph/store/package.json b/libs/@local/graph/store/package.json index 2fb5ede493d..04978d129c6 100644 --- a/libs/@local/graph/store/package.json +++ b/libs/@local/graph/store/package.json @@ -15,14 +15,11 @@ "./types": "./types/index.snap.js" }, "scripts": { - "build:codspeed": "cargo codspeed build -p hash-graph-store", "build:types": "INSTA_UPDATE=always mise exec --env dev cargo:cargo-insta -- cargo-insta test --features codegen --test codegen", "doc:dependency-diagram": "cargo run -p hash-repo-chores -- dependency-diagram --output docs/dependency-diagram.mmd --root hash-graph-store --root-deps-and-dependents --link-mode non-roots --include-dev-deps --include-build-deps --logging-console-level info", "fix:clippy": "just clippy --fix", "lint:clippy": "just clippy", "lint:tsc": "tsc --noEmit", - "test:codspeed": "cargo codspeed run -p hash-graph-store", - "test:miri": "cargo miri nextest run -- embedding::kernel embedding::clustering::tests::nearest_centroid_argmax_independent_of_inv_norm", "test:unit": "mise run test:unit @rust/hash-graph-store" }, "dependencies": { From eb38f87d283f629d2a76487c9ca022578133f591 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:29:55 +0200 Subject: [PATCH 27/38] chore: regenerate files --- apps/hash-graph/docs/dependency-diagram.mmd | 4 +- .../rust/docs/dependency-diagram.mmd | 32 ++-- libs/@local/codec/docs/dependency-diagram.mmd | 2 +- .../codegen/docs/dependency-diagram.mmd | 2 +- .../graph/api/docs/dependency-diagram.mmd | 4 +- .../authorization/docs/dependency-diagram.mmd | 32 ++-- .../embeddings/docs/dependency-diagram.mmd | 86 ++++++----- libs/@local/graph/embeddings/package.json | 3 + .../docs/dependency-diagram.mmd | 129 ++++++++-------- libs/@local/graph/postgres-store/package.json | 1 + .../graph/store/docs/dependency-diagram.mmd | 32 ++-- libs/@local/graph/store/package.json | 1 - .../docs/dependency-diagram.mmd | 2 +- .../type-fetcher/docs/dependency-diagram.mmd | 24 +-- .../graph/types/docs/dependency-diagram.mmd | 32 ++-- .../validation/docs/dependency-diagram.mmd | 30 ++-- .../harpc/server/docs/dependency-diagram.mmd | 26 ++-- .../harpc/types/docs/dependency-diagram.mmd | 2 +- .../wire-protocol/docs/dependency-diagram.mmd | 2 +- .../hashql/ast/docs/dependency-diagram.mmd | 141 +++++++++--------- .../compiletest/docs/dependency-diagram.mmd | 141 +++++++++--------- .../hashql/eval/docs/dependency-diagram.mmd | 141 +++++++++--------- .../hashql/hir/docs/dependency-diagram.mmd | 141 +++++++++--------- .../hashql/mir/docs/dependency-diagram.mmd | 141 +++++++++--------- .../syntax-jexpr/docs/dependency-diagram.mmd | 141 +++++++++--------- .../docs/dependency-diagram.mmd | 32 ++-- .../rust/docs/dependency-diagram.mmd | 34 ++--- 27 files changed, 664 insertions(+), 694 deletions(-) diff --git a/apps/hash-graph/docs/dependency-diagram.mmd b/apps/hash-graph/docs/dependency-diagram.mmd index ec8e7669a4d..c5aad87b71e 100644 --- a/apps/hash-graph/docs/dependency-diagram.mmd +++ b/apps/hash-graph/docs/dependency-diagram.mmd @@ -56,7 +56,6 @@ graph TD 1 -.-> 41 2 -.-> 3 2 --> 23 - 4 --> 6 4 --> 12 4 --> 13 4 --> 19 @@ -64,15 +63,16 @@ graph TD 4 --> 32 5 --> 1 6 --> 14 + 6 -.-> 37 7 --> 8 7 --> 34 + 9 --> 6 9 -.-> 7 9 --> 15 9 --> 33 10 --> 5 10 --> 14 10 --> 35 - 10 -.-> 37 11 --> 2 12 --> 33 13 --> 10 diff --git a/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd b/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd index 3cbfa718cbd..caefa648975 100644 --- a/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd +++ b/libs/@blockprotocol/type-system/rust/docs/dependency-diagram.mmd @@ -32,40 +32,35 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[darwin-kperf] - 24[darwin-kperf-criterion] - 25[darwin-kperf-events] - 26[darwin-kperf-sys] - 27[error-stack] - 28[hash-graph-benches] - 29[hash-graph-integration] - 30[hash-graph-test-data] + 23[error-stack] + 24[hash-graph-benches] + 25[hash-graph-integration] + 26[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 30 + 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 8 --> 22 - 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 30 - 12 -.-> 30 + 11 -.-> 26 + 12 -.-> 26 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 27 + 15 --> 23 16 -.-> 17 17 --> 18 17 --> 21 @@ -75,9 +70,6 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 23 --> 25 - 23 --> 26 - 24 --> 23 - 28 -.-> 4 - 29 -.-> 7 - 30 --> 8 + 24 -.-> 4 + 25 -.-> 7 + 26 --> 8 diff --git a/libs/@local/codec/docs/dependency-diagram.mmd b/libs/@local/codec/docs/dependency-diagram.mmd index 2dc72bb1081..046b44f0057 100644 --- a/libs/@local/codec/docs/dependency-diagram.mmd +++ b/libs/@local/codec/docs/dependency-diagram.mmd @@ -46,13 +46,13 @@ graph TD 1 -.-> 31 2 -.-> 3 2 --> 19 - 4 --> 6 4 --> 10 4 --> 15 4 --> 23 4 --> 26 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 diff --git a/libs/@local/codegen/docs/dependency-diagram.mmd b/libs/@local/codegen/docs/dependency-diagram.mmd index 46e354a3c6d..08c5bbadc20 100644 --- a/libs/@local/codegen/docs/dependency-diagram.mmd +++ b/libs/@local/codegen/docs/dependency-diagram.mmd @@ -42,13 +42,13 @@ graph TD 1 --> 9 1 -.-> 28 2 -.-> 3 - 4 --> 6 4 --> 10 4 --> 15 4 --> 21 4 --> 24 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 diff --git a/libs/@local/graph/api/docs/dependency-diagram.mmd b/libs/@local/graph/api/docs/dependency-diagram.mmd index ec3530323a2..b15023971bf 100644 --- a/libs/@local/graph/api/docs/dependency-diagram.mmd +++ b/libs/@local/graph/api/docs/dependency-diagram.mmd @@ -57,7 +57,6 @@ graph TD 1 -.-> 42 2 -.-> 3 2 --> 23 - 4 --> 6 4 --> 12 4 --> 13 4 --> 19 @@ -65,15 +64,16 @@ graph TD 4 --> 32 5 --> 1 6 --> 14 + 6 -.-> 37 7 --> 8 7 --> 34 + 9 --> 6 9 -.-> 7 9 --> 15 9 --> 33 10 --> 5 10 --> 14 10 --> 35 - 10 -.-> 37 11 --> 2 12 --> 33 13 --> 10 diff --git a/libs/@local/graph/authorization/docs/dependency-diagram.mmd b/libs/@local/graph/authorization/docs/dependency-diagram.mmd index 002c88588d2..aa9283d70f0 100644 --- a/libs/@local/graph/authorization/docs/dependency-diagram.mmd +++ b/libs/@local/graph/authorization/docs/dependency-diagram.mmd @@ -32,40 +32,35 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[darwin-kperf] - 24[darwin-kperf-criterion] - 25[darwin-kperf-events] - 26[darwin-kperf-sys] - 27[error-stack] - 28[hash-graph-benches] - 29[hash-graph-integration] - 30[hash-graph-test-data] + 23[error-stack] + 24[hash-graph-benches] + 25[hash-graph-integration] + 26[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 30 + 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 8 --> 22 - 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 30 - 12 -.-> 30 + 11 -.-> 26 + 12 -.-> 26 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 27 + 15 --> 23 16 -.-> 17 17 --> 18 17 --> 21 @@ -75,9 +70,6 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 23 --> 25 - 23 --> 26 - 24 --> 23 - 28 -.-> 4 - 29 -.-> 7 - 30 --> 8 + 24 -.-> 4 + 25 -.-> 7 + 26 --> 8 diff --git a/libs/@local/graph/embeddings/docs/dependency-diagram.mmd b/libs/@local/graph/embeddings/docs/dependency-diagram.mmd index 97ab6d97b76..5aa3d1e61bb 100644 --- a/libs/@local/graph/embeddings/docs/dependency-diagram.mmd +++ b/libs/@local/graph/embeddings/docs/dependency-diagram.mmd @@ -16,40 +16,58 @@ graph TD 5[hash-graph-authorization] 6[hash-graph-embeddings] class 6 root - 7[hash-graph-store] - 8[hash-graph-temporal-versioning] - 9[hash-graph-types] - 10[harpc-types] - 11[harpc-wire-protocol] - 12[hash-temporal-client] - 13[darwin-kperf] - 14[darwin-kperf-criterion] - 15[darwin-kperf-events] - 16[darwin-kperf-sys] - 17[error-stack] - 18[hash-graph-benches] - 19[hash-graph-test-data] + 7[hash-graph-postgres-store] + 8[hash-graph-store] + 9[hash-graph-temporal-versioning] + 10[hash-graph-types] + 11[harpc-types] + 12[harpc-wire-protocol] + 13[hashql-ast] + 14[hashql-compiletest] + 15[hashql-eval] + 16[hashql-hir] + 17[hashql-mir] + 18[hashql-syntax-jexpr] + 19[hash-temporal-client] + 20[darwin-kperf] + 21[darwin-kperf-criterion] + 22[darwin-kperf-events] + 23[darwin-kperf-sys] + 24[error-stack] + 25[hash-graph-benches] + 26[hash-graph-integration] + 27[hash-graph-test-data] 0 --> 4 - 1 --> 8 - 1 -.-> 19 + 1 --> 9 + 1 -.-> 27 2 -.-> 3 - 2 --> 11 - 4 --> 6 - 4 --> 7 + 2 --> 12 + 4 --> 15 + 4 --> 18 5 --> 1 - 6 --> 9 - 7 --> 5 - 7 --> 9 - 7 --> 12 - 7 -.-> 14 - 8 --> 2 - 9 -.-> 19 - 11 -.-> 10 - 11 --> 10 - 11 --> 17 - 12 --> 1 - 13 --> 15 - 13 --> 16 - 14 --> 13 - 18 -.-> 4 - 19 --> 7 + 6 --> 10 + 6 -.-> 21 + 7 --> 6 + 8 --> 5 + 8 --> 10 + 8 --> 19 + 9 --> 2 + 10 -.-> 27 + 12 -.-> 11 + 12 --> 11 + 12 --> 24 + 13 -.-> 14 + 14 --> 15 + 14 --> 18 + 15 --> 7 + 15 --> 17 + 16 -.-> 14 + 17 --> 16 + 18 --> 13 + 19 --> 1 + 20 --> 22 + 20 --> 23 + 21 --> 20 + 25 -.-> 4 + 26 -.-> 7 + 27 --> 8 diff --git a/libs/@local/graph/embeddings/package.json b/libs/@local/graph/embeddings/package.json index 15d7db0be60..61d00ad4ee1 100644 --- a/libs/@local/graph/embeddings/package.json +++ b/libs/@local/graph/embeddings/package.json @@ -15,5 +15,8 @@ "dependencies": { "@rust/error-stack": "workspace:*", "@rust/hash-graph-types": "workspace:*" + }, + "devDependencies": { + "@rust/darwin-kperf-criterion": "workspace:*" } } diff --git a/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd b/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd index d26e8f3bc11..eaa2eff809a 100644 --- a/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd +++ b/libs/@local/graph/postgres-store/docs/dependency-diagram.mmd @@ -14,69 +14,72 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - class 8 root - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - 16[hashql-compiletest] - 17[hashql-eval] - 18[hashql-hir] - 19[hashql-mir] - 20[hashql-syntax-jexpr] - 21[hash-status] - 22[hash-telemetry] - 23[hash-temporal-client] - 24[darwin-kperf] - 25[darwin-kperf-criterion] - 26[darwin-kperf-events] - 27[darwin-kperf-sys] - 28[error-stack] - 29[hash-graph-benches] - 30[hash-graph-integration] - 31[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + class 9 root + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + 17[hashql-compiletest] + 18[hashql-eval] + 19[hashql-hir] + 20[hashql-mir] + 21[hashql-syntax-jexpr] + 22[hash-status] + 23[hash-telemetry] + 24[hash-temporal-client] + 25[darwin-kperf] + 26[darwin-kperf-criterion] + 27[darwin-kperf-events] + 28[darwin-kperf-sys] + 29[error-stack] + 30[hash-graph-benches] + 31[hash-graph-integration] + 32[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 31 + 1 --> 11 + 1 -.-> 32 2 -.-> 3 - 2 --> 14 - 4 --> 17 - 4 --> 20 + 2 --> 15 + 4 --> 18 + 4 --> 21 5 --> 1 - 6 --> 7 - 6 --> 22 - 8 -.-> 6 - 8 --> 12 - 8 --> 21 - 9 --> 5 - 9 --> 11 - 9 --> 23 - 9 -.-> 25 - 10 --> 2 - 11 -.-> 31 - 12 -.-> 31 - 14 -.-> 13 - 14 --> 13 - 14 --> 28 - 15 -.-> 16 - 16 --> 17 - 16 --> 20 - 17 --> 8 - 17 --> 19 - 18 -.-> 16 - 19 --> 18 - 20 --> 15 - 22 --> 28 - 23 --> 1 - 24 --> 26 - 24 --> 27 - 25 --> 24 - 29 -.-> 4 - 30 -.-> 8 - 31 --> 9 + 6 --> 12 + 6 -.-> 26 + 7 --> 8 + 7 --> 23 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 22 + 10 --> 5 + 10 --> 12 + 10 --> 24 + 11 --> 2 + 12 -.-> 32 + 13 -.-> 32 + 15 -.-> 14 + 15 --> 14 + 15 --> 29 + 16 -.-> 17 + 17 --> 18 + 17 --> 21 + 18 --> 9 + 18 --> 20 + 19 -.-> 17 + 20 --> 19 + 21 --> 16 + 23 --> 29 + 24 --> 1 + 25 --> 27 + 25 --> 28 + 26 --> 25 + 30 -.-> 4 + 31 -.-> 9 + 32 --> 10 diff --git a/libs/@local/graph/postgres-store/package.json b/libs/@local/graph/postgres-store/package.json index 20b59a9c1a8..2ae76eade07 100644 --- a/libs/@local/graph/postgres-store/package.json +++ b/libs/@local/graph/postgres-store/package.json @@ -16,6 +16,7 @@ "@rust/error-stack": "workspace:*", "@rust/hash-codec": "workspace:*", "@rust/hash-graph-authorization": "workspace:*", + "@rust/hash-graph-embeddings": "workspace:*", "@rust/hash-graph-store": "workspace:*", "@rust/hash-graph-temporal-versioning": "workspace:*", "@rust/hash-graph-types": "workspace:*", diff --git a/libs/@local/graph/store/docs/dependency-diagram.mmd b/libs/@local/graph/store/docs/dependency-diagram.mmd index c25fc15d926..c4dd3693f57 100644 --- a/libs/@local/graph/store/docs/dependency-diagram.mmd +++ b/libs/@local/graph/store/docs/dependency-diagram.mmd @@ -32,40 +32,35 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[darwin-kperf] - 24[darwin-kperf-criterion] - 25[darwin-kperf-events] - 26[darwin-kperf-sys] - 27[error-stack] - 28[hash-graph-benches] - 29[hash-graph-integration] - 30[hash-graph-test-data] + 23[error-stack] + 24[hash-graph-benches] + 25[hash-graph-integration] + 26[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 30 + 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 8 --> 22 - 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 30 - 12 -.-> 30 + 11 -.-> 26 + 12 -.-> 26 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 27 + 15 --> 23 16 -.-> 17 17 --> 18 17 --> 21 @@ -75,9 +70,6 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 23 --> 25 - 23 --> 26 - 24 --> 23 - 28 -.-> 4 - 29 -.-> 7 - 30 --> 8 + 24 -.-> 4 + 25 -.-> 7 + 26 --> 8 diff --git a/libs/@local/graph/store/package.json b/libs/@local/graph/store/package.json index 04978d129c6..13a00b5532b 100644 --- a/libs/@local/graph/store/package.json +++ b/libs/@local/graph/store/package.json @@ -33,7 +33,6 @@ }, "devDependencies": { "@local/tsconfig": "workspace:*", - "@rust/darwin-kperf-criterion": "workspace:*", "@rust/hash-codegen": "workspace:*", "typescript": "5.9.3" } diff --git a/libs/@local/graph/temporal-versioning/docs/dependency-diagram.mmd b/libs/@local/graph/temporal-versioning/docs/dependency-diagram.mmd index 3df7e86b38a..a6593b53518 100644 --- a/libs/@local/graph/temporal-versioning/docs/dependency-diagram.mmd +++ b/libs/@local/graph/temporal-versioning/docs/dependency-diagram.mmd @@ -41,13 +41,13 @@ graph TD 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 diff --git a/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd b/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd index d1243620b21..5c3c10fd8c7 100644 --- a/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd +++ b/libs/@local/graph/type-fetcher/docs/dependency-diagram.mmd @@ -22,16 +22,12 @@ graph TD 10[harpc-types] 11[harpc-wire-protocol] 12[hash-temporal-client] - 13[darwin-kperf] - 14[darwin-kperf-criterion] - 15[darwin-kperf-events] - 16[darwin-kperf-sys] - 17[error-stack] - 18[hash-graph-benches] - 19[hash-graph-test-data] + 13[error-stack] + 14[hash-graph-benches] + 15[hash-graph-test-data] 0 --> 4 1 --> 7 - 1 -.-> 19 + 1 -.-> 15 2 -.-> 3 2 --> 11 4 --> 8 @@ -39,16 +35,12 @@ graph TD 6 --> 5 6 --> 9 6 --> 12 - 6 -.-> 14 7 --> 2 8 --> 6 - 9 -.-> 19 + 9 -.-> 15 11 -.-> 10 11 --> 10 - 11 --> 17 + 11 --> 13 12 --> 1 - 13 --> 15 - 13 --> 16 - 14 --> 13 - 18 -.-> 4 - 19 --> 6 + 14 -.-> 4 + 15 --> 6 diff --git a/libs/@local/graph/types/docs/dependency-diagram.mmd b/libs/@local/graph/types/docs/dependency-diagram.mmd index 1b887fb64f8..dfec2861ca5 100644 --- a/libs/@local/graph/types/docs/dependency-diagram.mmd +++ b/libs/@local/graph/types/docs/dependency-diagram.mmd @@ -32,40 +32,35 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[darwin-kperf] - 24[darwin-kperf-criterion] - 25[darwin-kperf-events] - 26[darwin-kperf-sys] - 27[error-stack] - 28[hash-graph-benches] - 29[hash-graph-integration] - 30[hash-graph-test-data] + 23[error-stack] + 24[hash-graph-benches] + 25[hash-graph-integration] + 26[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 30 + 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 8 --> 22 - 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 30 - 12 -.-> 30 + 11 -.-> 26 + 12 -.-> 26 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 27 + 15 --> 23 16 -.-> 17 17 --> 18 17 --> 21 @@ -75,9 +70,6 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 23 --> 25 - 23 --> 26 - 24 --> 23 - 28 -.-> 4 - 29 -.-> 7 - 30 --> 8 + 24 -.-> 4 + 25 -.-> 7 + 26 --> 8 diff --git a/libs/@local/graph/validation/docs/dependency-diagram.mmd b/libs/@local/graph/validation/docs/dependency-diagram.mmd index 4dd6fa9d4c7..a84b3c63ba9 100644 --- a/libs/@local/graph/validation/docs/dependency-diagram.mmd +++ b/libs/@local/graph/validation/docs/dependency-diagram.mmd @@ -29,17 +29,13 @@ graph TD 17[hashql-mir] 18[hashql-syntax-jexpr] 19[hash-temporal-client] - 20[darwin-kperf] - 21[darwin-kperf-criterion] - 22[darwin-kperf-events] - 23[darwin-kperf-sys] - 24[error-stack] - 25[hash-graph-benches] - 26[hash-graph-integration] - 27[hash-graph-test-data] + 20[error-stack] + 21[hash-graph-benches] + 22[hash-graph-integration] + 23[hash-graph-test-data] 0 --> 4 1 --> 8 - 1 -.-> 27 + 1 -.-> 23 2 -.-> 3 2 --> 12 4 --> 15 @@ -49,13 +45,12 @@ graph TD 7 --> 5 7 --> 9 7 --> 19 - 7 -.-> 21 8 --> 2 - 9 -.-> 27 - 10 -.-> 27 + 9 -.-> 23 + 10 -.-> 23 12 -.-> 11 12 --> 11 - 12 --> 24 + 12 --> 20 13 -.-> 14 14 --> 15 14 --> 18 @@ -65,9 +60,6 @@ graph TD 17 --> 16 18 --> 13 19 --> 1 - 20 --> 22 - 20 --> 23 - 21 --> 20 - 25 -.-> 4 - 26 -.-> 6 - 27 --> 7 + 21 -.-> 4 + 22 -.-> 6 + 23 --> 7 diff --git a/libs/@local/harpc/server/docs/dependency-diagram.mmd b/libs/@local/harpc/server/docs/dependency-diagram.mmd index 0f088c57070..c55bd6573e3 100644 --- a/libs/@local/harpc/server/docs/dependency-diagram.mmd +++ b/libs/@local/harpc/server/docs/dependency-diagram.mmd @@ -27,16 +27,12 @@ graph TD 15[harpc-types] 16[harpc-wire-protocol] 17[hash-temporal-client] - 18[darwin-kperf] - 19[darwin-kperf-criterion] - 20[darwin-kperf-events] - 21[darwin-kperf-sys] - 22[error-stack] - 23[hash-graph-benches] - 24[hash-graph-test-data] + 18[error-stack] + 19[hash-graph-benches] + 20[hash-graph-test-data] 0 --> 4 1 --> 7 - 1 -.-> 24 + 1 -.-> 20 2 -.-> 3 2 --> 16 4 --> 6 @@ -45,12 +41,11 @@ graph TD 6 --> 5 6 --> 8 6 --> 17 - 6 -.-> 19 7 --> 2 - 8 -.-> 24 + 8 -.-> 20 9 --> 13 10 --> 15 - 10 --> 22 + 10 --> 18 11 --> 2 11 -.-> 10 11 --> 10 @@ -61,10 +56,7 @@ graph TD 14 --> 11 16 -.-> 15 16 --> 15 - 16 --> 22 + 16 --> 18 17 --> 1 - 18 --> 20 - 18 --> 21 - 19 --> 18 - 23 -.-> 4 - 24 --> 6 + 19 -.-> 4 + 20 --> 6 diff --git a/libs/@local/harpc/types/docs/dependency-diagram.mmd b/libs/@local/harpc/types/docs/dependency-diagram.mmd index 91c4ae06387..e19512e8d9f 100644 --- a/libs/@local/harpc/types/docs/dependency-diagram.mmd +++ b/libs/@local/harpc/types/docs/dependency-diagram.mmd @@ -44,13 +44,13 @@ graph TD 1 --> 8 1 -.-> 30 2 --> 19 - 3 --> 5 3 --> 9 3 --> 15 3 --> 23 3 --> 26 4 --> 1 5 --> 10 + 6 --> 5 6 --> 11 7 --> 4 7 --> 10 diff --git a/libs/@local/harpc/wire-protocol/docs/dependency-diagram.mmd b/libs/@local/harpc/wire-protocol/docs/dependency-diagram.mmd index e665aaf4f40..4b9d3502f7b 100644 --- a/libs/@local/harpc/wire-protocol/docs/dependency-diagram.mmd +++ b/libs/@local/harpc/wire-protocol/docs/dependency-diagram.mmd @@ -44,13 +44,13 @@ graph TD 1 --> 8 1 -.-> 30 2 --> 18 - 3 --> 5 3 --> 9 3 --> 14 3 --> 22 3 --> 25 4 --> 1 5 --> 10 + 6 --> 5 6 --> 11 7 --> 4 7 --> 10 diff --git a/libs/@local/hashql/ast/docs/dependency-diagram.mmd b/libs/@local/hashql/ast/docs/dependency-diagram.mmd index 1d809d023df..49b505b3f35 100644 --- a/libs/@local/hashql/ast/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/ast/docs/dependency-diagram.mmd @@ -14,75 +14,78 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - class 15 root - 16[hashql-compiletest] - 17[hashql-core] - 18[hashql-diagnostics] - 19[hashql-eval] - 20[hashql-hir] - 21[hashql-macros] - 22[hashql-mir] - 23[hashql-syntax-jexpr] - 24[hash-status] - 25[hash-telemetry] - 26[hash-temporal-client] - 27[darwin-kperf] - 28[darwin-kperf-criterion] - 29[darwin-kperf-events] - 30[darwin-kperf-sys] - 31[error-stack] - 32[hash-graph-benches] - 33[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + class 16 root + 17[hashql-compiletest] + 18[hashql-core] + 19[hashql-diagnostics] + 20[hashql-eval] + 21[hashql-hir] + 22[hashql-macros] + 23[hashql-mir] + 24[hashql-syntax-jexpr] + 25[hash-status] + 26[hash-telemetry] + 27[hash-temporal-client] + 28[darwin-kperf] + 29[darwin-kperf-criterion] + 30[darwin-kperf-events] + 31[darwin-kperf-sys] + 32[error-stack] + 33[hash-graph-benches] + 34[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 33 + 1 --> 11 + 1 -.-> 34 2 -.-> 3 - 2 --> 14 - 4 --> 19 - 4 --> 23 + 2 --> 15 + 4 --> 20 + 4 --> 24 5 --> 1 - 6 --> 7 - 6 --> 25 - 8 -.-> 6 - 8 --> 12 - 8 --> 24 - 9 --> 5 - 9 --> 11 - 9 --> 26 - 9 -.-> 28 - 10 --> 2 - 11 -.-> 33 - 12 -.-> 33 - 14 -.-> 13 - 14 --> 13 - 14 --> 31 - 15 -.-> 16 - 16 --> 19 - 16 --> 23 - 17 --> 2 - 17 --> 18 - 17 --> 21 - 17 -.-> 28 - 19 --> 8 - 19 --> 22 - 20 -.-> 16 - 22 --> 20 - 23 --> 15 - 23 --> 17 - 25 --> 31 - 26 --> 1 - 27 --> 29 - 27 --> 30 - 28 --> 27 - 32 -.-> 4 - 33 --> 9 + 6 --> 12 + 6 -.-> 29 + 7 --> 8 + 7 --> 26 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 25 + 10 --> 5 + 10 --> 12 + 10 --> 27 + 11 --> 2 + 12 -.-> 34 + 13 -.-> 34 + 15 -.-> 14 + 15 --> 14 + 15 --> 32 + 16 -.-> 17 + 17 --> 20 + 17 --> 24 + 18 --> 2 + 18 --> 19 + 18 --> 22 + 18 -.-> 29 + 20 --> 9 + 20 --> 23 + 21 -.-> 17 + 23 --> 21 + 24 --> 16 + 24 --> 18 + 26 --> 32 + 27 --> 1 + 28 --> 30 + 28 --> 31 + 29 --> 28 + 33 -.-> 4 + 34 --> 10 diff --git a/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd b/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd index 7fa03c59383..9a2ec12d921 100644 --- a/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/compiletest/docs/dependency-diagram.mmd @@ -14,75 +14,78 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - 16[hashql-compiletest] - class 16 root - 17[hashql-core] - 18[hashql-diagnostics] - 19[hashql-eval] - 20[hashql-hir] - 21[hashql-macros] - 22[hashql-mir] - 23[hashql-syntax-jexpr] - 24[hash-status] - 25[hash-telemetry] - 26[hash-temporal-client] - 27[darwin-kperf] - 28[darwin-kperf-criterion] - 29[darwin-kperf-events] - 30[darwin-kperf-sys] - 31[error-stack] - 32[hash-graph-benches] - 33[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + 17[hashql-compiletest] + class 17 root + 18[hashql-core] + 19[hashql-diagnostics] + 20[hashql-eval] + 21[hashql-hir] + 22[hashql-macros] + 23[hashql-mir] + 24[hashql-syntax-jexpr] + 25[hash-status] + 26[hash-telemetry] + 27[hash-temporal-client] + 28[darwin-kperf] + 29[darwin-kperf-criterion] + 30[darwin-kperf-events] + 31[darwin-kperf-sys] + 32[error-stack] + 33[hash-graph-benches] + 34[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 33 + 1 --> 11 + 1 -.-> 34 2 -.-> 3 - 2 --> 14 - 4 --> 19 - 4 --> 23 + 2 --> 15 + 4 --> 20 + 4 --> 24 5 --> 1 - 6 --> 7 - 6 --> 25 - 8 -.-> 6 - 8 --> 12 - 8 --> 24 - 9 --> 5 - 9 --> 11 - 9 --> 26 - 9 -.-> 28 - 10 --> 2 - 11 -.-> 33 - 12 -.-> 33 - 14 -.-> 13 - 14 --> 13 - 14 --> 31 - 15 -.-> 16 - 16 --> 19 - 16 --> 23 - 17 --> 2 - 17 --> 18 - 17 --> 21 - 17 -.-> 28 - 19 --> 8 - 19 --> 22 - 20 -.-> 16 - 22 --> 20 - 23 --> 15 - 23 --> 17 - 25 --> 31 - 26 --> 1 - 27 --> 29 - 27 --> 30 - 28 --> 27 - 32 -.-> 4 - 33 --> 9 + 6 --> 12 + 6 -.-> 29 + 7 --> 8 + 7 --> 26 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 25 + 10 --> 5 + 10 --> 12 + 10 --> 27 + 11 --> 2 + 12 -.-> 34 + 13 -.-> 34 + 15 -.-> 14 + 15 --> 14 + 15 --> 32 + 16 -.-> 17 + 17 --> 20 + 17 --> 24 + 18 --> 2 + 18 --> 19 + 18 --> 22 + 18 -.-> 29 + 20 --> 9 + 20 --> 23 + 21 -.-> 17 + 23 --> 21 + 24 --> 16 + 24 --> 18 + 26 --> 32 + 27 --> 1 + 28 --> 30 + 28 --> 31 + 29 --> 28 + 33 -.-> 4 + 34 --> 10 diff --git a/libs/@local/hashql/eval/docs/dependency-diagram.mmd b/libs/@local/hashql/eval/docs/dependency-diagram.mmd index ba720211915..0817d44f171 100644 --- a/libs/@local/hashql/eval/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/eval/docs/dependency-diagram.mmd @@ -14,75 +14,78 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - 16[hashql-compiletest] - 17[hashql-core] - 18[hashql-diagnostics] - 19[hashql-eval] - class 19 root - 20[hashql-hir] - 21[hashql-macros] - 22[hashql-mir] - 23[hashql-syntax-jexpr] - 24[hash-status] - 25[hash-telemetry] - 26[hash-temporal-client] - 27[darwin-kperf] - 28[darwin-kperf-criterion] - 29[darwin-kperf-events] - 30[darwin-kperf-sys] - 31[error-stack] - 32[hash-graph-benches] - 33[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + 17[hashql-compiletest] + 18[hashql-core] + 19[hashql-diagnostics] + 20[hashql-eval] + class 20 root + 21[hashql-hir] + 22[hashql-macros] + 23[hashql-mir] + 24[hashql-syntax-jexpr] + 25[hash-status] + 26[hash-telemetry] + 27[hash-temporal-client] + 28[darwin-kperf] + 29[darwin-kperf-criterion] + 30[darwin-kperf-events] + 31[darwin-kperf-sys] + 32[error-stack] + 33[hash-graph-benches] + 34[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 33 + 1 --> 11 + 1 -.-> 34 2 -.-> 3 - 2 --> 14 - 4 --> 19 - 4 --> 23 + 2 --> 15 + 4 --> 20 + 4 --> 24 5 --> 1 - 6 --> 7 - 6 --> 25 - 8 -.-> 6 - 8 --> 12 - 8 --> 24 - 9 --> 5 - 9 --> 11 - 9 --> 26 - 9 -.-> 28 - 10 --> 2 - 11 -.-> 33 - 12 -.-> 33 - 14 -.-> 13 - 14 --> 13 - 14 --> 31 - 15 -.-> 16 - 16 --> 19 - 16 --> 23 - 17 --> 2 - 17 --> 18 - 17 --> 21 - 17 -.-> 28 - 19 --> 8 - 19 --> 22 - 20 -.-> 16 - 22 --> 20 - 23 --> 15 - 23 --> 17 - 25 --> 31 - 26 --> 1 - 27 --> 29 - 27 --> 30 - 28 --> 27 - 32 -.-> 4 - 33 --> 9 + 6 --> 12 + 6 -.-> 29 + 7 --> 8 + 7 --> 26 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 25 + 10 --> 5 + 10 --> 12 + 10 --> 27 + 11 --> 2 + 12 -.-> 34 + 13 -.-> 34 + 15 -.-> 14 + 15 --> 14 + 15 --> 32 + 16 -.-> 17 + 17 --> 20 + 17 --> 24 + 18 --> 2 + 18 --> 19 + 18 --> 22 + 18 -.-> 29 + 20 --> 9 + 20 --> 23 + 21 -.-> 17 + 23 --> 21 + 24 --> 16 + 24 --> 18 + 26 --> 32 + 27 --> 1 + 28 --> 30 + 28 --> 31 + 29 --> 28 + 33 -.-> 4 + 34 --> 10 diff --git a/libs/@local/hashql/hir/docs/dependency-diagram.mmd b/libs/@local/hashql/hir/docs/dependency-diagram.mmd index 31c175f2525..d0f49158af5 100644 --- a/libs/@local/hashql/hir/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/hir/docs/dependency-diagram.mmd @@ -14,75 +14,78 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - 16[hashql-compiletest] - 17[hashql-core] - 18[hashql-diagnostics] - 19[hashql-eval] - 20[hashql-hir] - class 20 root - 21[hashql-macros] - 22[hashql-mir] - 23[hashql-syntax-jexpr] - 24[hash-status] - 25[hash-telemetry] - 26[hash-temporal-client] - 27[darwin-kperf] - 28[darwin-kperf-criterion] - 29[darwin-kperf-events] - 30[darwin-kperf-sys] - 31[error-stack] - 32[hash-graph-benches] - 33[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + 17[hashql-compiletest] + 18[hashql-core] + 19[hashql-diagnostics] + 20[hashql-eval] + 21[hashql-hir] + class 21 root + 22[hashql-macros] + 23[hashql-mir] + 24[hashql-syntax-jexpr] + 25[hash-status] + 26[hash-telemetry] + 27[hash-temporal-client] + 28[darwin-kperf] + 29[darwin-kperf-criterion] + 30[darwin-kperf-events] + 31[darwin-kperf-sys] + 32[error-stack] + 33[hash-graph-benches] + 34[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 33 + 1 --> 11 + 1 -.-> 34 2 -.-> 3 - 2 --> 14 - 4 --> 19 - 4 --> 23 + 2 --> 15 + 4 --> 20 + 4 --> 24 5 --> 1 - 6 --> 7 - 6 --> 25 - 8 -.-> 6 - 8 --> 12 - 8 --> 24 - 9 --> 5 - 9 --> 11 - 9 --> 26 - 9 -.-> 28 - 10 --> 2 - 11 -.-> 33 - 12 -.-> 33 - 14 -.-> 13 - 14 --> 13 - 14 --> 31 - 15 -.-> 16 - 16 --> 19 - 16 --> 23 - 17 --> 2 - 17 --> 18 - 17 --> 21 - 17 -.-> 28 - 19 --> 8 - 19 --> 22 - 20 -.-> 16 - 22 --> 20 - 23 --> 15 - 23 --> 17 - 25 --> 31 - 26 --> 1 - 27 --> 29 - 27 --> 30 - 28 --> 27 - 32 -.-> 4 - 33 --> 9 + 6 --> 12 + 6 -.-> 29 + 7 --> 8 + 7 --> 26 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 25 + 10 --> 5 + 10 --> 12 + 10 --> 27 + 11 --> 2 + 12 -.-> 34 + 13 -.-> 34 + 15 -.-> 14 + 15 --> 14 + 15 --> 32 + 16 -.-> 17 + 17 --> 20 + 17 --> 24 + 18 --> 2 + 18 --> 19 + 18 --> 22 + 18 -.-> 29 + 20 --> 9 + 20 --> 23 + 21 -.-> 17 + 23 --> 21 + 24 --> 16 + 24 --> 18 + 26 --> 32 + 27 --> 1 + 28 --> 30 + 28 --> 31 + 29 --> 28 + 33 -.-> 4 + 34 --> 10 diff --git a/libs/@local/hashql/mir/docs/dependency-diagram.mmd b/libs/@local/hashql/mir/docs/dependency-diagram.mmd index 351c3ebd1f6..08cc3a23141 100644 --- a/libs/@local/hashql/mir/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/mir/docs/dependency-diagram.mmd @@ -14,75 +14,78 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - 16[hashql-compiletest] - 17[hashql-core] - 18[hashql-diagnostics] - 19[hashql-eval] - 20[hashql-hir] - 21[hashql-macros] - 22[hashql-mir] - class 22 root - 23[hashql-syntax-jexpr] - 24[hash-status] - 25[hash-telemetry] - 26[hash-temporal-client] - 27[darwin-kperf] - 28[darwin-kperf-criterion] - 29[darwin-kperf-events] - 30[darwin-kperf-sys] - 31[error-stack] - 32[hash-graph-benches] - 33[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + 17[hashql-compiletest] + 18[hashql-core] + 19[hashql-diagnostics] + 20[hashql-eval] + 21[hashql-hir] + 22[hashql-macros] + 23[hashql-mir] + class 23 root + 24[hashql-syntax-jexpr] + 25[hash-status] + 26[hash-telemetry] + 27[hash-temporal-client] + 28[darwin-kperf] + 29[darwin-kperf-criterion] + 30[darwin-kperf-events] + 31[darwin-kperf-sys] + 32[error-stack] + 33[hash-graph-benches] + 34[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 33 + 1 --> 11 + 1 -.-> 34 2 -.-> 3 - 2 --> 14 - 4 --> 19 - 4 --> 23 + 2 --> 15 + 4 --> 20 + 4 --> 24 5 --> 1 - 6 --> 7 - 6 --> 25 - 8 -.-> 6 - 8 --> 12 - 8 --> 24 - 9 --> 5 - 9 --> 11 - 9 --> 26 - 9 -.-> 28 - 10 --> 2 - 11 -.-> 33 - 12 -.-> 33 - 14 -.-> 13 - 14 --> 13 - 14 --> 31 - 15 -.-> 16 - 16 --> 19 - 16 --> 23 - 17 --> 2 - 17 --> 18 - 17 --> 21 - 17 -.-> 28 - 19 --> 8 - 19 --> 22 - 20 -.-> 16 - 22 --> 20 - 23 --> 15 - 23 --> 17 - 25 --> 31 - 26 --> 1 - 27 --> 29 - 27 --> 30 - 28 --> 27 - 32 -.-> 4 - 33 --> 9 + 6 --> 12 + 6 -.-> 29 + 7 --> 8 + 7 --> 26 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 25 + 10 --> 5 + 10 --> 12 + 10 --> 27 + 11 --> 2 + 12 -.-> 34 + 13 -.-> 34 + 15 -.-> 14 + 15 --> 14 + 15 --> 32 + 16 -.-> 17 + 17 --> 20 + 17 --> 24 + 18 --> 2 + 18 --> 19 + 18 --> 22 + 18 -.-> 29 + 20 --> 9 + 20 --> 23 + 21 -.-> 17 + 23 --> 21 + 24 --> 16 + 24 --> 18 + 26 --> 32 + 27 --> 1 + 28 --> 30 + 28 --> 31 + 29 --> 28 + 33 -.-> 4 + 34 --> 10 diff --git a/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd b/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd index 1abbd2f75d5..659be40ce02 100644 --- a/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd +++ b/libs/@local/hashql/syntax-jexpr/docs/dependency-diagram.mmd @@ -14,75 +14,78 @@ graph TD 3[hash-codegen] 4[hash-graph-api] 5[hash-graph-authorization] - 6[hash-graph-migrations] - 7[hash-graph-migrations-macros] - 8[hash-graph-postgres-store] - 9[hash-graph-store] - 10[hash-graph-temporal-versioning] - 11[hash-graph-types] - 12[hash-graph-validation] - 13[harpc-types] - 14[harpc-wire-protocol] - 15[hashql-ast] - 16[hashql-compiletest] - 17[hashql-core] - 18[hashql-diagnostics] - 19[hashql-eval] - 20[hashql-hir] - 21[hashql-macros] - 22[hashql-mir] - 23[hashql-syntax-jexpr] - class 23 root - 24[hash-status] - 25[hash-telemetry] - 26[hash-temporal-client] - 27[darwin-kperf] - 28[darwin-kperf-criterion] - 29[darwin-kperf-events] - 30[darwin-kperf-sys] - 31[error-stack] - 32[hash-graph-benches] - 33[hash-graph-test-data] + 6[hash-graph-embeddings] + 7[hash-graph-migrations] + 8[hash-graph-migrations-macros] + 9[hash-graph-postgres-store] + 10[hash-graph-store] + 11[hash-graph-temporal-versioning] + 12[hash-graph-types] + 13[hash-graph-validation] + 14[harpc-types] + 15[harpc-wire-protocol] + 16[hashql-ast] + 17[hashql-compiletest] + 18[hashql-core] + 19[hashql-diagnostics] + 20[hashql-eval] + 21[hashql-hir] + 22[hashql-macros] + 23[hashql-mir] + 24[hashql-syntax-jexpr] + class 24 root + 25[hash-status] + 26[hash-telemetry] + 27[hash-temporal-client] + 28[darwin-kperf] + 29[darwin-kperf-criterion] + 30[darwin-kperf-events] + 31[darwin-kperf-sys] + 32[error-stack] + 33[hash-graph-benches] + 34[hash-graph-test-data] 0 --> 4 - 1 --> 10 - 1 -.-> 33 + 1 --> 11 + 1 -.-> 34 2 -.-> 3 - 2 --> 14 - 4 --> 19 - 4 --> 23 + 2 --> 15 + 4 --> 20 + 4 --> 24 5 --> 1 - 6 --> 7 - 6 --> 25 - 8 -.-> 6 - 8 --> 12 - 8 --> 24 - 9 --> 5 - 9 --> 11 - 9 --> 26 - 9 -.-> 28 - 10 --> 2 - 11 -.-> 33 - 12 -.-> 33 - 14 -.-> 13 - 14 --> 13 - 14 --> 31 - 15 -.-> 16 - 16 --> 19 - 16 --> 23 - 17 --> 2 - 17 --> 18 - 17 --> 21 - 17 -.-> 28 - 19 --> 8 - 19 --> 22 - 20 -.-> 16 - 22 --> 20 - 23 --> 15 - 23 --> 17 - 25 --> 31 - 26 --> 1 - 27 --> 29 - 27 --> 30 - 28 --> 27 - 32 -.-> 4 - 33 --> 9 + 6 --> 12 + 6 -.-> 29 + 7 --> 8 + 7 --> 26 + 9 --> 6 + 9 -.-> 7 + 9 --> 13 + 9 --> 25 + 10 --> 5 + 10 --> 12 + 10 --> 27 + 11 --> 2 + 12 -.-> 34 + 13 -.-> 34 + 15 -.-> 14 + 15 --> 14 + 15 --> 32 + 16 -.-> 17 + 17 --> 20 + 17 --> 24 + 18 --> 2 + 18 --> 19 + 18 --> 22 + 18 -.-> 29 + 20 --> 9 + 20 --> 23 + 21 -.-> 17 + 23 --> 21 + 24 --> 16 + 24 --> 18 + 26 --> 32 + 27 --> 1 + 28 --> 30 + 28 --> 31 + 29 --> 28 + 33 -.-> 4 + 34 --> 10 diff --git a/libs/@local/temporal-client/docs/dependency-diagram.mmd b/libs/@local/temporal-client/docs/dependency-diagram.mmd index 1152c9ec8ed..3a4af027246 100644 --- a/libs/@local/temporal-client/docs/dependency-diagram.mmd +++ b/libs/@local/temporal-client/docs/dependency-diagram.mmd @@ -32,40 +32,35 @@ graph TD 21[hashql-syntax-jexpr] 22[hash-temporal-client] class 22 root - 23[darwin-kperf] - 24[darwin-kperf-criterion] - 25[darwin-kperf-events] - 26[darwin-kperf-sys] - 27[error-stack] - 28[hash-graph-benches] - 29[hash-graph-integration] - 30[hash-graph-test-data] + 23[error-stack] + 24[hash-graph-benches] + 25[hash-graph-integration] + 26[hash-graph-test-data] 0 --> 4 1 --> 9 - 1 -.-> 30 + 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 8 --> 22 - 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 30 - 12 -.-> 30 + 11 -.-> 26 + 12 -.-> 26 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 27 + 15 --> 23 16 -.-> 17 17 --> 18 17 --> 21 @@ -75,9 +70,6 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 23 --> 25 - 23 --> 26 - 24 --> 23 - 28 -.-> 4 - 29 -.-> 7 - 30 --> 8 + 24 -.-> 4 + 25 -.-> 7 + 26 --> 8 diff --git a/tests/graph/test-data/rust/docs/dependency-diagram.mmd b/tests/graph/test-data/rust/docs/dependency-diagram.mmd index fc41eb6288a..0390665dcd2 100644 --- a/tests/graph/test-data/rust/docs/dependency-diagram.mmd +++ b/tests/graph/test-data/rust/docs/dependency-diagram.mmd @@ -31,41 +31,36 @@ graph TD 20[hashql-mir] 21[hashql-syntax-jexpr] 22[hash-temporal-client] - 23[darwin-kperf] - 24[darwin-kperf-criterion] - 25[darwin-kperf-events] - 26[darwin-kperf-sys] - 27[error-stack] - 28[hash-graph-benches] - 29[hash-graph-integration] - 30[hash-graph-test-data] - class 30 root + 23[error-stack] + 24[hash-graph-benches] + 25[hash-graph-integration] + 26[hash-graph-test-data] + class 26 root 0 --> 4 1 --> 9 - 1 -.-> 30 + 1 -.-> 26 2 -.-> 3 2 --> 15 - 4 --> 6 4 --> 10 4 --> 13 4 --> 18 4 --> 21 5 --> 1 6 --> 11 + 7 --> 6 7 --> 12 8 --> 5 8 --> 11 8 --> 22 - 8 -.-> 24 9 --> 2 10 --> 8 - 11 -.-> 30 - 12 -.-> 30 + 11 -.-> 26 + 12 -.-> 26 13 -.-> 1 13 --> 14 15 -.-> 14 15 --> 14 - 15 --> 27 + 15 --> 23 16 -.-> 17 17 --> 18 17 --> 21 @@ -75,9 +70,6 @@ graph TD 20 --> 19 21 --> 16 22 --> 1 - 23 --> 25 - 23 --> 26 - 24 --> 23 - 28 -.-> 4 - 29 -.-> 7 - 30 --> 8 + 24 -.-> 4 + 25 -.-> 7 + 26 --> 8 From 5f7fd585ab9ecaf7f9ed956c321a0baa14a2f634 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:36:58 +0200 Subject: [PATCH 28/38] fix: the darn yarn lockfile --- yarn.lock | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/yarn.lock b/yarn.lock index cb4ccc74b26..230a410d724 100644 --- a/yarn.lock +++ b/yarn.lock @@ -13141,6 +13141,7 @@ __metadata: version: 0.0.0-use.local resolution: "@rust/hash-graph-embeddings@workspace:libs/@local/graph/embeddings" dependencies: + "@rust/darwin-kperf-criterion": "workspace:*" "@rust/error-stack": "workspace:*" "@rust/hash-graph-types": "workspace:*" languageName: unknown @@ -13197,6 +13198,7 @@ __metadata: "@rust/error-stack": "workspace:*" "@rust/hash-codec": "workspace:*" "@rust/hash-graph-authorization": "workspace:*" + "@rust/hash-graph-embeddings": "workspace:*" "@rust/hash-graph-migrations": "workspace:*" "@rust/hash-graph-store": "workspace:*" "@rust/hash-graph-temporal-versioning": "workspace:*" @@ -13215,7 +13217,6 @@ __metadata: dependencies: "@blockprotocol/type-system-rs": "workspace:*" "@local/tsconfig": "workspace:*" - "@rust/darwin-kperf-criterion": "workspace:*" "@rust/error-stack": "workspace:*" "@rust/hash-codec": "workspace:*" "@rust/hash-codegen": "workspace:*" From 77ada17179e3dc6119e0a191c5bf0a14c1a751ad Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:47:04 +0200 Subject: [PATCH 29/38] chore: remove tautological tests --- libs/@local/graph/embeddings/src/dimension.rs | 39 ------------------- 1 file changed, 39 deletions(-) diff --git a/libs/@local/graph/embeddings/src/dimension.rs b/libs/@local/graph/embeddings/src/dimension.rs index 3701ee85a91..cf4f021ea82 100644 --- a/libs/@local/graph/embeddings/src/dimension.rs +++ b/libs/@local/graph/embeddings/src/dimension.rs @@ -50,42 +50,3 @@ pub const D256: Dimension = Dimension(NonZero::new(256).unwrap()); pub const D512: Dimension = Dimension(NonZero::new(512).unwrap()); pub const D1536: Dimension = Dimension(NonZero::new(1536).unwrap()); pub const D3072: Dimension = Dimension(NonZero::new(3072).unwrap()); - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn valid_multiples_of_8() { - for value in [8, 16, 24, 128, 256, 3072] { - assert!( - Dimension::new(value).is_some(), - "{value} should be a valid dimension" - ); - } - } - - #[test] - fn zero_rejected() { - assert!(Dimension::new(0).is_none()); - } - - #[test] - fn non_multiples_of_8_rejected() { - for value in [1, 2, 3, 4, 5, 6, 7, 9, 10, 15, 17, 100, 3071] { - assert!( - Dimension::new(value).is_none(), - "{value} should not be a valid dimension" - ); - } - } - - #[test] - fn constants_have_correct_values() { - assert_eq!(D128.0.get(), 128); - assert_eq!(D256.0.get(), 256); - assert_eq!(D512.0.get(), 512); - assert_eq!(D1536.0.get(), 1536); - assert_eq!(D3072.0.get(), 3072); - } -} From 02d295824ad8ae0103ce13fa4c3cbdd115b287c9 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:57:05 +0200 Subject: [PATCH 30/38] feat: only initialize once. --- apps/hash-graph/src/subcommand/mod.rs | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/apps/hash-graph/src/subcommand/mod.rs b/apps/hash-graph/src/subcommand/mod.rs index e6e7ca8e334..dbb31a7289e 100644 --- a/apps/hash-graph/src/subcommand/mod.rs +++ b/apps/hash-graph/src/subcommand/mod.rs @@ -7,7 +7,7 @@ mod snapshot; mod type_fetcher; use core::time::Duration; -use std::{thread::available_parallelism, time::Instant}; +use std::{sync::Once, thread::available_parallelism, time::Instant}; use clap::Parser; use error_stack::{Report, ensure}; @@ -147,15 +147,18 @@ fn block_on( service_name: &'static str, tracing_config: TracingConfig, ) -> Result<(), Report> { - rayon::ThreadPoolBuilder::new() - .num_threads( - available_parallelism() - .map_or(1, |cores| cores.get() / 2) - .max(1), - ) - .thread_name(|index| format!("rayon-{index}")) - .build_global() - .expect("rayon pool should be initialized exactly once"); + static THREAD_POOL: Once = Once::new(); + THREAD_POOL.call_once(|| { + rayon::ThreadPoolBuilder::new() + .num_threads( + available_parallelism() + .map_or(1, |cores| cores.get() / 2) + .max(1), + ) + .thread_name(|index| format!("rayon-{index}")) + .build_global() + .expect("rayon pool should be initialized exactly once"); + }); tokio::runtime::Builder::new_multi_thread() .enable_all() From 1d4a1f7c6e68de1b2ad3e0ba3b39efacf38cb0c1 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:41:50 +0200 Subject: [PATCH 31/38] fix: docs --- libs/@local/graph/embeddings/benches/clustering.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/@local/graph/embeddings/benches/clustering.rs b/libs/@local/graph/embeddings/benches/clustering.rs index 24658df6723..6cd7f86a5be 100644 --- a/libs/@local/graph/embeddings/benches/clustering.rs +++ b/libs/@local/graph/embeddings/benches/clustering.rs @@ -9,7 +9,7 @@ //! spread across the rayon pool and per-thread instruction counts would only see the calling //! thread. //! -//! [`cluster`]: hash_graph_store::embedding::clustering::cluster +//! [`cluster`]: hash_graph_embeddings::clustering::cluster #![expect( unsafe_code, clippy::float_arithmetic, From 0bb8726f5bb2ea9cb915cc1d38e61adc564a8271 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:10:18 +0200 Subject: [PATCH 32/38] feat: move to rayon spawns --- Cargo.lock | 1 + libs/@local/graph/postgres-store/Cargo.toml | 1 + .../store/postgres/knowledge/entity/mod.rs | 94 +++++++++++++++++-- 3 files changed, 87 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c4912b58916..b3659f34603 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3789,6 +3789,7 @@ dependencies = [ "indoc", "postgres-types", "pretty_assertions", + "rayon", "refinery", "regex", "semver", diff --git a/libs/@local/graph/postgres-store/Cargo.toml b/libs/@local/graph/postgres-store/Cargo.toml index 114c71642d4..af076b8ab31 100644 --- a/libs/@local/graph/postgres-store/Cargo.toml +++ b/libs/@local/graph/postgres-store/Cargo.toml @@ -40,6 +40,7 @@ derive_more = { workspace = true } dotenv-flow = { workspace = true } futures = { workspace = true } postgres-types = { workspace = true, features = ["derive", "with-serde_json-1"] } +rayon = { workspace = true } refinery = { workspace = true, features = ["tokio-postgres"] } regex = { workspace = true } semver = { workspace = true, features = ["serde"] } diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 719c23e8d00..87629031197 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -4,7 +4,7 @@ mod read; mod summary; use alloc::{borrow::Cow, collections::BTreeMap}; -use core::{borrow::Borrow as _, mem}; +use core::{any::Any, borrow::Borrow as _, fmt, mem}; use std::collections::{HashMap, HashSet}; use error_stack::{FutureExt as _, Report, ResultExt as _, TryReportStreamExt as _, ensure}; @@ -65,6 +65,7 @@ use hash_graph_types::{ use hash_graph_validation::{EntityPreprocessor, Validate as _}; use hash_status::StatusCode; use postgres_types::ToSql; +use tokio::sync::oneshot; use tokio_postgres::{GenericClient as _, error::SqlState}; use tracing::Instrument as _; use type_system::{ @@ -113,6 +114,71 @@ use crate::store::{ validation::StoreProvider, }; +/// The panic that happened during a spawned task. +/// +/// Opaque to fulfil the `Sync` contract, which has the safety requirement that it must be sound for +/// `&JoinError`, to cross thread boundaries. By design, a `&JoinError` has no API whatsoever, +/// making it useless, thus harmless, thus memory safe. +/// +/// This use has precedent, see the nightly `SyncView`, the `SyncWrapper` inside tokio, and +/// `SyncWrapper` of the `sync_wrapper` crate. +struct JoinError(Box); + +// SAFETY: An immutable reference to a `JoinError` is useless, as the value can only be interacted +// with by getting the inner value. This mirrors the design of `SyncView`, see the rationale behind +// it. We choose to implement a custom wrapper instead, to be able to downcast, as long as the +// actual value hidden behind is `Sync`, making the wrapper a no-op, mirroring the internal +// `SyncWrapper` type of tokio, used for it's `JoinError`. +// See: https://github.com/tokio-rs/tokio/blob/c4c6265a0746a79d4a2f3852f726aa0101f29fd3/tokio/src/util/sync_wrapper.rs#L8 +// and: https://github.com/rust-lang/rust/blob/f10db292a3733b5c67c8da8c7661195ff4b05774/library/core/src/sync/sync_view.rs#L90 +#[expect(unsafe_code)] +unsafe impl Sync for JoinError {} + +impl JoinError { + // Adapted from: https://github.com/rust-lang/rust/blob/6c8138de8f1c96b2f66adbbc0e37c73525444750/library/std/src/panicking.rs#L779-L787 + fn message(&self) -> Option<&str> { + if let Some(value) = self.downcast_ref_sync::<&'static str>() { + return Some(*value); + } + + if let Some(value) = self.downcast_ref_sync::() { + return Some(&**value); + } + + None + } + + fn downcast_ref_sync(&self) -> Option<&T> { + // If the downcast fails, the inner value is not touched, so no thread-safety violation can + // occur. + self.0.downcast_ref() + } +} + +impl fmt::Debug for JoinError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut debug = fmt.debug_tuple("JoinError"); + + if let Some(message) = self.message() { + return debug.field(&message).finish(); + } + + debug.finish_non_exhaustive() + } +} + +impl fmt::Display for JoinError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + if let Some(message) = self.message() { + return write!(fmt, "task panicked with message: {message}"); + } + + fmt.write_str("task panicked") + } +} + +impl core::error::Error for JoinError {} + impl PostgresStore where C: AsClient, @@ -2756,15 +2822,25 @@ where }), ); - let result = tokio::task::spawn_blocking(move || { - hash_graph_embeddings::clustering::cluster(&flat, dimension, &config) - }) - .await - .change_context(ClusterError::Store)?; + let (tx, rx) = oneshot::channel(); + rayon::spawn(move || { + let result = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| { + hash_graph_embeddings::clustering::cluster(&flat, dimension, &config) + })); + + let result = result.map_err(JoinError); + + let _tx = tx.send(result); + }); + + let clustering = rx + .await + .change_context(ClusterError::Store)? + .change_context(ClusterError::Store)?; let mut groups: BTreeMap> = BTreeMap::new(); for (index, id) in found_ids.iter().enumerate() { - groups.entry(result.label(index)).or_default().push(*id); + groups.entry(clustering.label(index)).or_default().push(*id); } let clusters = groups @@ -2772,14 +2848,14 @@ where .map(|(cluster_id, entity_ids)| EntityCluster { cluster_id, entity_ids, - centroid: result.centroid(cluster_id).to_vec(), + centroid: clustering.centroid(cluster_id).to_vec(), }) .collect(); Ok(ClusterEntitiesResponse { clusters, missing_embeddings, - inertia: result.inertia, + inertia: clustering.inertia, }) } } From 4119460ec8b7e29797a800059ce17b5a13895170 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:28:42 +0200 Subject: [PATCH 33/38] feat: test --- Cargo.lock | 1 + .../store/postgres/knowledge/entity/mod.rs | 3 +- tests/graph/integration/Cargo.toml | 1 + tests/graph/integration/package.json | 1 + .../graph/integration/postgres/clustering.rs | 377 ++++++++++++++++++ tests/graph/integration/postgres/lib.rs | 1 + 6 files changed, 383 insertions(+), 1 deletion(-) create mode 100644 tests/graph/integration/postgres/clustering.rs diff --git a/Cargo.lock b/Cargo.lock index b3659f34603..1503da65e31 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3706,6 +3706,7 @@ dependencies = [ "hash-graph-store", "hash-graph-temporal-versioning", "hash-graph-test-data", + "hash-graph-types", "hash-status", "hash-telemetry", "pretty_assertions", diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 87629031197..8b877f27f49 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -59,6 +59,7 @@ use hash_graph_temporal_versioning::{ TransactionTime, }; use hash_graph_types::{ + Embedding, knowledge::property::visitor::EntityVisitor as _, ontology::{DataTypeLookup, OntologyTypeProvider}, }; @@ -89,7 +90,7 @@ use type_system::{ entity_type::{ClosedEntityType, ClosedMultiEntityType, EntityTypeUuid}, id::{OntologyTypeUuid, VersionedUrl}, }, - principal::actor::ActorEntityUuid, + principal::{actor::ActorEntityUuid, actor_group::WebId}, }; use uuid::Uuid; diff --git a/tests/graph/integration/Cargo.toml b/tests/graph/integration/Cargo.toml index 4e289457f43..da1d6be507f 100644 --- a/tests/graph/integration/Cargo.toml +++ b/tests/graph/integration/Cargo.toml @@ -15,6 +15,7 @@ hash-graph-postgres-store = { workspace = true } hash-graph-store = { workspace = true } hash-graph-temporal-versioning = { workspace = true } hash-graph-test-data = { workspace = true } +hash-graph-types = { workspace = true } hash-status = { workspace = true } hash-telemetry = { workspace = true } type-system = { workspace = true } diff --git a/tests/graph/integration/package.json b/tests/graph/integration/package.json index 8fd6b60ede6..0fdebc36434 100644 --- a/tests/graph/integration/package.json +++ b/tests/graph/integration/package.json @@ -18,6 +18,7 @@ "@rust/hash-graph-store": "workspace:*", "@rust/hash-graph-temporal-versioning": "workspace:*", "@rust/hash-graph-test-data": "workspace:*", + "@rust/hash-graph-types": "workspace:*", "@rust/hash-status": "workspace:*", "@rust/hash-telemetry": "workspace:*" } diff --git a/tests/graph/integration/postgres/clustering.rs b/tests/graph/integration/postgres/clustering.rs new file mode 100644 index 00000000000..a7016fe4229 --- /dev/null +++ b/tests/graph/integration/postgres/clustering.rs @@ -0,0 +1,377 @@ +use core::num::NonZero; +use std::collections::HashSet; + +use hash_graph_store::{ + entity::{ + ClusterEntitiesParams, CreateEntityParams, EntityStore as _, UpdateEntityEmbeddingsParams, + }, + error::ClusterError, +}; +use hash_graph_temporal_versioning::Timestamp; +use hash_graph_test_data::{data_type, entity, entity_type, property_type}; +use hash_graph_types::{Embedding, knowledge::entity::EntityEmbedding}; +use type_system::{ + knowledge::{ + entity::{EntityId, id::EntityUuid, provenance::ProvidedEntityEditionProvenance}, + property::{PropertyObject, PropertyObjectWithMetadata}, + }, + ontology::id::{BaseUrl, OntologyTypeVersion, VersionedUrl}, + principal::{actor::ActorType, actor_group::WebId}, + provenance::{OriginProvenance, OriginType}, +}; +use uuid::Uuid; + +use crate::{DatabaseApi, DatabaseTestWrapper}; + +/// Dimension used for clustering requests in these tests. +/// +/// Embeddings are stored as 3072-dimensional vectors but matryoshka-truncated +/// server-side, so only the first `CLUSTER_DIM` components carry signal here. +const CLUSTER_DIM: u16 = 8; + +async fn seed(database: &mut DatabaseTestWrapper) -> DatabaseApi<'_> { + database + .seed( + [ + data_type::VALUE_V1, + data_type::TEXT_V1, + data_type::NUMBER_V1, + ], + [ + property_type::NAME_V1, + property_type::AGE_V1, + property_type::FAVORITE_SONG_V1, + property_type::FAVORITE_FILM_V1, + property_type::HOBBY_V1, + property_type::INTERESTS_V1, + ], + [ + entity_type::LINK_V1, + entity_type::link::FRIEND_OF_V1, + entity_type::link::ACQUAINTANCE_OF_V1, + entity_type::PERSON_V1, + ], + ) + .await + .expect("could not seed database") +} + +fn person_entity_type_id() -> VersionedUrl { + VersionedUrl { + base_url: BaseUrl::new( + "https://blockprotocol.org/@alice/types/entity-type/person/".to_owned(), + ) + .expect("couldn't construct Base URL"), + version: OntologyTypeVersion { + major: 1, + pre_release: None, + }, + } +} + +async fn create_person(api: &mut DatabaseApi<'_>) -> EntityId { + let person: PropertyObject = + serde_json::from_str(entity::PERSON_ALICE_V1).expect("could not parse entity"); + + api.create_entity( + api.account_id, + CreateEntityParams { + web_id: WebId::new(api.account_id), + entity_uuid: None, + decision_time: None, + entity_type_ids: HashSet::from([person_entity_type_id()]), + properties: PropertyObjectWithMetadata::from_parts(person, None) + .expect("could not create property with metadata object"), + confidence: None, + link_data: None, + draft: false, + policies: Vec::new(), + provenance: ProvidedEntityEditionProvenance { + actor_type: ActorType::User, + origin: OriginProvenance::from_empty_type(OriginType::Api), + sources: Vec::new(), + }, + read_only: false, + }, + ) + .await + .expect("could not create entity") + .metadata + .record_id + .entity_id +} + +/// Builds a full-width stored embedding pointing along `axis` within the first +/// [`CLUSTER_DIM`] components, with a small per-entity `jitter` so vectors in +/// the same group are distinct but remain tightly clustered. +#[expect(clippy::indexing_slicing, clippy::float_arithmetic)] +fn embedding_along_axis(axis: usize, jitter_axis: usize, jitter: f32) -> Embedding<'static> { + assert!(axis < usize::from(CLUSTER_DIM)); + assert!(jitter_axis < usize::from(CLUSTER_DIM)); + + let mut vector = vec![0.0_f32; Embedding::DIM]; + vector[axis] = 1.0; + vector[jitter_axis] += jitter; + Embedding::from(vector) +} + +async fn insert_embedding( + api: &mut DatabaseApi<'_>, + entity_id: EntityId, + embedding: Embedding<'static>, +) { + api.update_entity_embeddings( + api.account_id, + UpdateEntityEmbeddingsParams { + entity_id, + embeddings: vec![EntityEmbedding { + property: None, + embedding, + }], + updated_at_transaction_time: Timestamp::now(), + updated_at_decision_time: Timestamp::now(), + reset: false, + }, + ) + .await + .expect("could not insert entity embedding"); +} + +const fn cluster_params(entity_ids: Vec, cluster_count: u16) -> ClusterEntitiesParams { + ClusterEntitiesParams { + entity_ids, + cluster_count, + dimension: NonZero::new(CLUSTER_DIM).expect("dimension should be non-zero"), + seed: Some(0), + } +} + +#[tokio::test] +async fn clusters_entities_by_embedding_direction() { + let mut database = DatabaseTestWrapper::new().await; + let mut api = seed(&mut database).await; + + // Two well-separated groups: group A points along axis 0, group B along + // axis 1. Spherical k-means with `k = 2` must separate them. + let mut group_a = Vec::new(); + for index in 0..3_usize { + let entity_id = create_person(&mut api).await; + #[expect(clippy::cast_precision_loss, reason = "tiny test-only jitter values")] + insert_embedding( + &mut api, + entity_id, + embedding_along_axis(0, 2, 0.01 * index as f32), + ) + .await; + group_a.push(entity_id); + } + + let mut group_b = Vec::new(); + for index in 0..2_usize { + let entity_id = create_person(&mut api).await; + #[expect(clippy::cast_precision_loss, reason = "tiny test-only jitter values")] + insert_embedding( + &mut api, + entity_id, + embedding_along_axis(1, 3, 0.01 * index as f32), + ) + .await; + group_b.push(entity_id); + } + + // One existing entity without any stored embedding ... + let entity_without_embedding = create_person(&mut api).await; + // ... and one entity ID that does not exist at all. + let nonexistent_entity = EntityId { + web_id: WebId::new(api.account_id), + entity_uuid: EntityUuid::new(Uuid::new_v4()), + draft_id: None, + }; + + let mut requested = group_a.clone(); + requested.extend(&group_b); + requested.push(entity_without_embedding); + requested.push(nonexistent_entity); + + let response = api + .cluster_entities(api.account_id, cluster_params(requested, 2)) + .await + .expect("could not cluster entities"); + + assert_eq!( + response + .missing_embeddings + .iter() + .copied() + .collect::>(), + HashSet::from([entity_without_embedding, nonexistent_entity]), + "entities without embeddings and unknown entities should be reported as missing" + ); + + assert_eq!(response.clusters.len(), 2, "expected exactly two clusters"); + + let group_a_set: HashSet = group_a.iter().copied().collect(); + let group_b_set: HashSet = group_b.iter().copied().collect(); + + let cluster_with_a = response + .clusters + .iter() + .find(|cluster| cluster.entity_ids.contains(&group_a[0])) + .expect("group A should be assigned to a cluster"); + let cluster_with_b = response + .clusters + .iter() + .find(|cluster| cluster.entity_ids.contains(&group_b[0])) + .expect("group B should be assigned to a cluster"); + + assert_eq!( + cluster_with_a + .entity_ids + .iter() + .copied() + .collect::>(), + group_a_set, + "group A should form one cluster" + ); + assert_eq!( + cluster_with_b + .entity_ids + .iter() + .copied() + .collect::>(), + group_b_set, + "group B should form the other cluster" + ); + + for cluster in &response.clusters { + assert_eq!( + cluster.centroid.len(), + usize::from(CLUSTER_DIM), + "centroid length should match the requested dimension" + ); + } + + assert!( + response.inertia < 0.01, + "tightly clustered groups should have near-zero inertia, got {}", + response.inertia + ); +} + +#[tokio::test] +async fn permission_denied_is_reported_as_missing() { + let mut database = DatabaseTestWrapper::new().await; + let mut api = seed(&mut database).await; + + let entity_id = create_person(&mut api).await; + insert_embedding(&mut api, entity_id, embedding_along_axis(0, 1, 0.0)).await; + + // The owning actor can cluster the entity. + let response = api + .cluster_entities(api.account_id, cluster_params(vec![entity_id], 1)) + .await + .expect("could not cluster entities"); + assert_eq!(response.clusters.len(), 1); + assert!(response.missing_embeddings.is_empty()); + + // A machine actor without access to the web must not see the entity. To + // avoid leaking permission information, the entity is reported exactly as + // if it had no embedding. + let machine_id = api.create_machine("clustering-outsider").await; + let response = api + .cluster_entities(machine_id.into(), cluster_params(vec![entity_id], 1)) + .await + .expect("could not cluster entities"); + + assert!( + response.clusters.is_empty(), + "unauthorized entities should not be clustered" + ); + assert_eq!( + response.missing_embeddings, + vec![entity_id], + "unauthorized entities should be indistinguishable from missing embeddings" + ); + assert!(response.inertia.abs() < f32::EPSILON); +} + +#[tokio::test] +async fn zero_cluster_count_returns_no_clusters() { + let mut database = DatabaseTestWrapper::new().await; + let mut api = seed(&mut database).await; + + let entity_id = create_person(&mut api).await; + insert_embedding(&mut api, entity_id, embedding_along_axis(0, 1, 0.0)).await; + + let response = api + .cluster_entities(api.account_id, cluster_params(vec![entity_id], 0)) + .await + .expect("could not cluster entities"); + + assert!(response.clusters.is_empty()); + assert!( + response.missing_embeddings.is_empty(), + "entities with embeddings should not be reported as missing even when `k = 0`" + ); + assert!(response.inertia.abs() < f32::EPSILON); +} + +#[tokio::test] +async fn rejects_invalid_parameters() { + let mut database = DatabaseTestWrapper::new().await; + let api = seed(&mut database).await; + + // Dimension must be a multiple of 8. + let report = api + .cluster_entities( + api.account_id, + ClusterEntitiesParams { + entity_ids: Vec::new(), + cluster_count: 2, + dimension: NonZero::new(7).expect("dimension should be non-zero"), + seed: Some(0), + }, + ) + .await + .expect_err("dimension not a multiple of 8 should be rejected"); + assert!(matches!( + report.current_context(), + ClusterError::InvalidDimension { .. } + )); + + // Dimension must not exceed 512. + let report = api + .cluster_entities( + api.account_id, + ClusterEntitiesParams { + entity_ids: Vec::new(), + cluster_count: 2, + dimension: NonZero::new(520).expect("dimension should be non-zero"), + seed: Some(0), + }, + ) + .await + .expect_err("dimension above 512 should be rejected"); + assert!(matches!( + report.current_context(), + ClusterError::DimensionTooLarge { .. } + )); + + // Cluster count must not exceed 64. + let report = api + .cluster_entities( + api.account_id, + ClusterEntitiesParams { + entity_ids: Vec::new(), + cluster_count: 65, + dimension: NonZero::new(CLUSTER_DIM).expect("dimension should be non-zero"), + seed: Some(0), + }, + ) + .await + .expect_err("cluster count above 64 should be rejected"); + assert!(matches!( + report.current_context(), + ClusterError::KTooLarge { .. } + )); +} diff --git a/tests/graph/integration/postgres/lib.rs b/tests/graph/integration/postgres/lib.rs index f9e54f423bb..59053b59abb 100644 --- a/tests/graph/integration/postgres/lib.rs +++ b/tests/graph/integration/postgres/lib.rs @@ -6,6 +6,7 @@ extern crate alloc; +mod clustering; mod data_type; mod drafts; mod email_filter_protection; From 3fd0938fca241d36c788f588c3673b37314604f1 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:30:31 +0200 Subject: [PATCH 34/38] chore: lockfile --- yarn.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/yarn.lock b/yarn.lock index 230a410d724..9e2bfaa85e3 100644 --- a/yarn.lock +++ b/yarn.lock @@ -13169,6 +13169,7 @@ __metadata: "@rust/hash-graph-store": "workspace:*" "@rust/hash-graph-temporal-versioning": "workspace:*" "@rust/hash-graph-test-data": "workspace:*" + "@rust/hash-graph-types": "workspace:*" "@rust/hash-status": "workspace:*" "@rust/hash-telemetry": "workspace:*" languageName: unknown From 464e3b0392289aa1bfa5c99bf12d411c7ae9fff4 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:40:39 +0200 Subject: [PATCH 35/38] fix: suggestions from (external) code review --- apps/hash-graph/src/args.rs | 15 +- apps/hash-graph/src/main.rs | 3 +- apps/hash-graph/src/subcommand/mod.rs | 164 ++++++++++++++---- apps/hash-graph/src/subcommand/server.rs | 12 +- libs/@local/graph/api/openapi/openapi.json | 3 +- libs/@local/graph/api/src/rest/entity/mod.rs | 28 ++- libs/@local/graph/api/src/rest/mod.rs | 5 +- .../@local/graph/embeddings/src/clustering.rs | 1 + libs/@local/graph/postgres-store/src/lib.rs | 2 +- .../store/postgres/knowledge/entity/mod.rs | 58 +++---- libs/@local/graph/store/src/entity/store.rs | 2 +- .../graph/integration/postgres/clustering.rs | 2 +- 12 files changed, 221 insertions(+), 74 deletions(-) diff --git a/apps/hash-graph/src/args.rs b/apps/hash-graph/src/args.rs index b3c538276a3..32b717f1316 100644 --- a/apps/hash-graph/src/args.rs +++ b/apps/hash-graph/src/args.rs @@ -7,7 +7,7 @@ use clap::{ }; use hash_telemetry::TracingConfig; -use crate::subcommand::Subcommand; +use crate::subcommand::{Subcommand, WorkerThreads}; /// Arguments passed to the program. #[derive(Debug, Parser)] @@ -16,6 +16,19 @@ pub struct Args { #[clap(flatten)] pub tracing_config: TracingConfig, + /// Number of threads in the global worker pool used for CPU-bound work such as entity + /// clustering. + /// + /// Accepts a fixed count (e.g. `4`) or a count relative to the available CPU cores: `n` for + /// all cores, `n/2` for half, `n/4` for a quarter, and so on. + #[clap( + long, + global = true, + default_value_t, + env = "HASH_GRAPH_WORKER_THREADS" + )] + pub worker_threads: WorkerThreads, + /// Specify a subcommand to run. #[command(subcommand)] pub subcommand: Subcommand, diff --git a/apps/hash-graph/src/main.rs b/apps/hash-graph/src/main.rs index e8ac0681ffc..efdcb62d0b9 100644 --- a/apps/hash-graph/src/main.rs +++ b/apps/hash-graph/src/main.rs @@ -30,9 +30,10 @@ fn main() -> Result<(), Report> { let Args { subcommand, tracing_config, + worker_threads, } = Args::parse_args(); let _sentry_guard = init(&tracing_config.sentry, release_name!()); - subcommand.execute(tracing_config) + subcommand.execute(tracing_config, worker_threads) } diff --git a/apps/hash-graph/src/subcommand/mod.rs b/apps/hash-graph/src/subcommand/mod.rs index dbb31a7289e..d4a4e24d1a2 100644 --- a/apps/hash-graph/src/subcommand/mod.rs +++ b/apps/hash-graph/src/subcommand/mod.rs @@ -6,7 +6,7 @@ mod server; mod snapshot; mod type_fetcher; -use core::time::Duration; +use core::{fmt, num::NonZero, str::FromStr, time::Duration}; use std::{sync::Once, thread::available_parallelism, time::Instant}; use clap::Parser; @@ -15,6 +15,19 @@ use hash_telemetry::{TracingConfig, init_tracing}; use tokio::time::sleep; use tokio_util::{sync::CancellationToken, task::TaskTracker}; +pub use self::{ + admin_server::{AdminServerArgs, admin_server}, + completions::{CompletionsArgs, completions}, + migrate::{MigrateArgs, migrate}, + server::{ServerArgs, server}, + snapshot::{SnapshotArgs, snapshot}, + type_fetcher::{TypeFetcherArgs, type_fetcher}, +}; +use crate::{ + error::{GraphError, HealthcheckError}, + subcommand::reindex_cache::{ReindexCacheArgs, reindex_cache}, +}; + /// Drop guard that fires the `abort` token when a server task exits unexpectedly. /// /// "Unexpectedly" means the `shutdown` token has not been cancelled yet. This covers both @@ -87,6 +100,85 @@ impl ServerLifecycle { } } +/// Number of threads for the global worker pool used for CPU-bound work. +/// +/// Parses either a fixed thread count (e.g. `4`) or a count relative to the number of available +/// CPU cores: `n` for all cores, `n/2` for half of them, `n/4` for a quarter, and so on. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum WorkerThreads { + /// The available CPU cores divided by the given divisor (`n`, `n/2`, `n/4`, ...). + Cores { divisor: NonZero }, + /// A fixed number of threads. + Fixed(NonZero), +} + +impl WorkerThreads { + /// Resolves to a concrete thread count, clamped to at least one thread. + #[expect( + clippy::integer_division, + reason = "Deriving a thread count from the core count is inherently lossy." + )] + fn resolve(self) -> NonZero { + match self { + Self::Fixed(threads) => threads, + Self::Cores { divisor } => available_parallelism() + .ok() + .and_then(|cores| NonZero::new(cores.get() / divisor)) + .unwrap_or(NonZero::::MIN), + } + } +} + +impl Default for WorkerThreads { + fn default() -> Self { + const HALF: NonZero = NonZero::new(2).expect("two should be non-zero"); + Self::Cores { divisor: HALF } + } +} + +impl fmt::Display for WorkerThreads { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + Self::Cores { divisor } if divisor == NonZero::::MIN => fmt.write_str("n"), + Self::Cores { divisor } => write!(fmt, "n/{divisor}"), + Self::Fixed(threads) => write!(fmt, "{threads}"), + } + } +} + +/// Error returned when parsing a [`WorkerThreads`] value fails. +#[derive(Debug)] +pub struct ParseWorkerThreadsError; + +impl fmt::Display for ParseWorkerThreadsError { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.write_str("expected a positive integer, `n`, or `n/` (e.g. `4`, `n`, `n/2`)") + } +} + +impl core::error::Error for ParseWorkerThreadsError {} + +impl FromStr for WorkerThreads { + type Err = ParseWorkerThreadsError; + + fn from_str(value: &str) -> Result { + match value.strip_prefix(['n', 'N']) { + Some("") => Ok(Self::Cores { + divisor: NonZero::::MIN, + }), + Some(rest) => rest + .strip_prefix('/') + .and_then(|divisor| divisor.parse().ok()) + .map(|divisor| Self::Cores { divisor }) + .ok_or(ParseWorkerThreadsError), + None => value + .parse() + .map(Self::Fixed) + .map_err(|_error: core::num::ParseIntError| ParseWorkerThreadsError), + } + } +} + /// Shared healthcheck arguments for all server subcommands. #[derive(Debug, Clone, Parser)] pub(crate) struct HealthcheckArgs { @@ -103,19 +195,6 @@ pub(crate) struct HealthcheckArgs { pub timeout: Option, } -pub use self::{ - admin_server::{AdminServerArgs, admin_server}, - completions::{CompletionsArgs, completions}, - migrate::{MigrateArgs, migrate}, - server::{ServerArgs, server}, - snapshot::{SnapshotArgs, snapshot}, - type_fetcher::{TypeFetcherArgs, type_fetcher}, -}; -use crate::{ - error::{GraphError, HealthcheckError}, - subcommand::reindex_cache::{ReindexCacheArgs, reindex_cache}, -}; - /// Subcommand for the program. #[derive(Debug, clap::Subcommand)] pub enum Subcommand { @@ -141,20 +220,16 @@ pub enum Subcommand { ReindexCache(Box), } -#[expect(clippy::integer_division, clippy::integer_division_remainder_used)] fn block_on( future: impl Future>>, service_name: &'static str, tracing_config: TracingConfig, + worker_threads: WorkerThreads, ) -> Result<(), Report> { static THREAD_POOL: Once = Once::new(); THREAD_POOL.call_once(|| { rayon::ThreadPoolBuilder::new() - .num_threads( - available_parallelism() - .map_or(1, |cores| cores.get() / 2) - .max(1), - ) + .num_threads(worker_threads.resolve().get()) .thread_name(|index| format!("rayon-{index}")) .build_global() .expect("rayon pool should be initialized exactly once"); @@ -173,24 +248,49 @@ fn block_on( } impl Subcommand { - pub(crate) fn execute(self, tracing_config: TracingConfig) -> Result<(), Report> { + pub(crate) fn execute( + self, + tracing_config: TracingConfig, + worker_threads: WorkerThreads, + ) -> Result<(), Report> { match self { - Self::Server(args) => block_on(server(*args), "Graph API", tracing_config), - Self::AdminServer(args) => { - block_on(admin_server(*args), "Graph Admin API", tracing_config) - } - Self::Migrate(args) => block_on(migrate(*args), "Graph Migrations", tracing_config), - Self::TypeFetcher(args) => { - block_on(type_fetcher(*args), "Type Fetcher", tracing_config) + Self::Server(args) => { + block_on(server(*args), "Graph API", tracing_config, worker_threads) } + Self::AdminServer(args) => block_on( + admin_server(*args), + "Graph Admin API", + tracing_config, + worker_threads, + ), + Self::Migrate(args) => block_on( + migrate(*args), + "Graph Migrations", + tracing_config, + worker_threads, + ), + Self::TypeFetcher(args) => block_on( + type_fetcher(*args), + "Type Fetcher", + tracing_config, + worker_threads, + ), Self::Completions(ref args) => { completions(args); Ok(()) } - Self::Snapshot(args) => block_on(snapshot(*args), "Graph Snapshot", tracing_config), - Self::ReindexCache(args) => { - block_on(reindex_cache(*args), "Graph Indexer", tracing_config) - } + Self::Snapshot(args) => block_on( + snapshot(*args), + "Graph Snapshot", + tracing_config, + worker_threads, + ), + Self::ReindexCache(args) => block_on( + reindex_cache(*args), + "Graph Indexer", + tracing_config, + worker_threads, + ), } } } diff --git a/apps/hash-graph/src/subcommand/server.rs b/apps/hash-graph/src/subcommand/server.rs index 8a9bd213701..47922c55b43 100644 --- a/apps/hash-graph/src/subcommand/server.rs +++ b/apps/hash-graph/src/subcommand/server.rs @@ -16,8 +16,8 @@ use harpc_server::Server; use hash_codec::bytes::JsonLinesEncoder; use hash_graph_api::{ rest::{ - ApiConfig, QueryLogger, RestApiStore, RestRouterDependencies, hashql::CompilerContext, - rest_api_router, + ApiConfig, QueryLogger, RestApiStore, RestRouterDependencies, entity::ClusteringContext, + hashql::CompilerContext, rest_api_router, }, rpc::Dependencies, }; @@ -239,6 +239,13 @@ pub struct ServerConfig { #[clap(flatten)] pub compiler: CompilerConfig, + /// Maximum number of entity-clustering requests processed at the same time. + /// + /// Excess requests wait until a slot frees up. If not set, the number of concurrent + /// clustering requests is unbounded. + #[clap(long, env = "HASH_GRAPH_CLUSTERING_CONCURRENCY_LIMIT")] + pub clustering_concurrency_limit: Option>, + /// Outputs the queries made to the graph to the specified file. #[clap(long)] pub log_queries: Option, @@ -457,6 +464,7 @@ where query_logger, api_config: config.api_config, compiler, + clustering: Arc::new(ClusteringContext::new(config.clustering_concurrency_limit)), }); start_rest_server(router, config.http_address, lifecycle); diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index 332ce2efcaa..cc55cd23940 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -3856,7 +3856,8 @@ "items": { "$ref": "#/components/schemas/EntityId" }, - "description": "Entities from the request that had no stored embedding or that do not exist." + "description": "Entities from the request that had no stored embedding or that do not exist.", + "uniqueItems": true } } }, diff --git a/libs/@local/graph/api/src/rest/entity/mod.rs b/libs/@local/graph/api/src/rest/entity/mod.rs index afc026cf3f1..bc10f42d808 100644 --- a/libs/@local/graph/api/src/rest/entity/mod.rs +++ b/libs/@local/graph/api/src/rest/entity/mod.rs @@ -3,10 +3,12 @@ pub mod query; use alloc::sync::Arc; +use core::num::NonZero; use std::collections::HashMap; use axum::{Extension, Router, routing::post}; use error_stack::{Report, ResultExt as _}; +use futures::future::OptionFuture; use hash_graph_authorization::policies::principal::actor::AuthenticatedActor; use hash_graph_embeddings::OpenAiEmbeddingClient; use hash_graph_postgres_store::store::error::{EntityDoesNotExist, RaceConditionOnUpdate}; @@ -45,6 +47,7 @@ use hash_graph_types::{ }; use hash_temporal_client::TemporalClient; use serde::Deserialize as _; +use tokio::sync::Semaphore; use type_system::{ knowledge::{ Confidence, Entity, Property, @@ -606,6 +609,19 @@ where .map_err(report_to_response) } +pub struct ClusteringContext { + pub limit: Option, +} + +impl ClusteringContext { + #[must_use] + pub fn new(concurrency_limit: Option>) -> Self { + Self { + limit: concurrency_limit.map(|limit| Semaphore::new(limit.get())), + } + } +} + #[utoipa::path( post, path = "/entities/embeddings/clusters", @@ -623,15 +639,21 @@ where )] async fn cluster_entities( AuthenticatedUserHeader(actor_id): AuthenticatedUserHeader, - store_pool: Extension>, - temporal_client: Extension>>, + Extension(store_pool): Extension>, + Extension(temporal_client): Extension>>, + Extension(context): Extension>, Json(params): Json, ) -> Result, BoxedResponse> where S: StorePool + Send + Sync, { + let _permit = OptionFuture::from(context.limit.as_ref().map(Semaphore::acquire)) + .await + .transpose() + .expect("semaphore should never be closed"); + let store = store_pool - .acquire(temporal_client.0) + .acquire(temporal_client) .await .map_err(report_to_response)?; diff --git a/libs/@local/graph/api/src/rest/mod.rs b/libs/@local/graph/api/src/rest/mod.rs index 3ebd2a169d9..00e6751be6d 100644 --- a/libs/@local/graph/api/src/rest/mod.rs +++ b/libs/@local/graph/api/src/rest/mod.rs @@ -100,6 +100,7 @@ use utoipa::{ use uuid::Uuid; use self::{ + entity::ClusteringContext, status::{BoxedResponse, report_to_response, status_to_response}, utoipa_typedef::{ MaybeListOfDataTypeMetadata, MaybeListOfEntityTypeMetadata, @@ -543,6 +544,7 @@ where pub query_logger: Option, pub api_config: ApiConfig, pub compiler: Arc, + pub clustering: Arc, } /// A [`Router`] that only serves the `OpenAPI` specification (JSON, and necessary subschemas) for @@ -593,7 +595,8 @@ where .layer(Extension(dependencies.embedding_client)) .layer(Extension(dependencies.domain_regex)) .layer(Extension(dependencies.api_config)) - .layer(Extension(dependencies.compiler)); + .layer(Extension(dependencies.compiler)) + .layer(Extension(Arc::new(dependencies.clustering))); if let Some(query_logger) = dependencies.query_logger { router = router.layer(Extension(query_logger)); diff --git a/libs/@local/graph/embeddings/src/clustering.rs b/libs/@local/graph/embeddings/src/clustering.rs index 45111c4b8f6..0837494cf08 100644 --- a/libs/@local/graph/embeddings/src/clustering.rs +++ b/libs/@local/graph/embeddings/src/clustering.rs @@ -27,6 +27,7 @@ use super::{dimension::Dimension, kernel}; /// /// Use [`Config::for_k_with_seed`] to construct with reasonable defaults, then override individual /// fields as needed. +#[derive(Debug, Copy, Clone)] pub struct Config { /// Number of clusters. pub k: u16, diff --git a/libs/@local/graph/postgres-store/src/lib.rs b/libs/@local/graph/postgres-store/src/lib.rs index 4fc05125e87..7733674896f 100644 --- a/libs/@local/graph/postgres-store/src/lib.rs +++ b/libs/@local/graph/postgres-store/src/lib.rs @@ -9,7 +9,7 @@ // Library Features extend_one, - iter_intersperse, + iter_intersperse )] #![cfg_attr(not(miri), doc(test(attr(deny(warnings, clippy::all)))))] #![expect( diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 8b877f27f49..4a2eca1817c 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -3,7 +3,7 @@ mod query; mod read; mod summary; -use alloc::{borrow::Cow, collections::BTreeMap}; +use alloc::borrow::Cow; use core::{any::Any, borrow::Borrow as _, fmt, mem}; use std::collections::{HashMap, HashSet}; @@ -17,7 +17,7 @@ use hash_graph_authorization::policies::{ resource::{EntityResourceConstraint, ResourceConstraint}, store::{PolicyCreationParams, PrincipalStore as _}, }; -use hash_graph_embeddings::Dimension; +use hash_graph_embeddings::{Dimension, clustering::Clustering}; use hash_graph_store::{ entity::{ ClusterEntitiesParams, ClusterEntitiesResponse, CreateEntityParams, DeleteEntitiesParams, @@ -2673,9 +2673,9 @@ where actor_id: ActorEntityUuid, params: ClusterEntitiesParams, ) -> Result> { - const { assert!(Embedding::DIM <= u16::MAX as usize) }; const MAX_ALLOWED_DIM: u16 = 512; const MAX_ALLOWED_K: u16 = 64; + const { assert!(Embedding::DIM <= u16::MAX as usize) }; let dimension = Dimension::new(params.dimension.get()).ok_or_else(|| { Report::new(ClusterError::InvalidDimension { @@ -2719,15 +2719,15 @@ where .await .change_context(ClusterError::Store)?; - let permitted_ids: Vec = params + let permitted_ids: Vec<_> = params .entity_ids .iter() .filter(|&id| permitted.contains_key(id)) .copied() .collect(); - let entity_uuids: Vec = permitted_ids.iter().map(|id| id.entity_uuid).collect(); - let web_ids: Vec = permitted_ids.iter().map(|id| id.web_id).collect(); + let entity_uuids: Vec<_> = permitted_ids.iter().map(|id| id.entity_uuid).collect(); + let web_ids: Vec<_> = permitted_ids.iter().map(|id| id.web_id).collect(); // Truncate server-side via `subvector` so postgres only sends // `truncated_dim`-dimensional vectors over the wire. @@ -2765,8 +2765,12 @@ where let mut row_stream = core::pin::pin!(row_stream); - let mut flat: Vec = Vec::with_capacity(permitted_ids.len() * truncated_dim); - let mut found_ids: Vec = Vec::with_capacity(permitted_ids.len()); + let mut flat: Vec<_> = Vec::with_capacity(permitted_ids.len() * truncated_dim); + let mut found_ids: Vec<_> = Vec::with_capacity(permitted_ids.len()); + + // Every requested entity not in a cluster goes into `missing_embeddings`, whether due to + // permissions or no embedding. Distinguishing the two would leak permission information. + let mut missing_ids: HashSet<_> = params.entity_ids.iter().copied().collect(); while let Some(row) = row_stream .try_next() @@ -2779,30 +2783,19 @@ where flat.extend(embedding.iter()); - found_ids.push(EntityId { + let id = EntityId { web_id, entity_uuid, draft_id: None, - }); + }; + found_ids.push(id); + missing_ids.remove(&id); } - // Every requested entity not in a cluster goes into - // `missing_embeddings`, whether due to permissions or no embedding. - // Distinguishing the two would leak permission information. - let found_set: HashSet<(WebId, EntityUuid)> = found_ids - .iter() - .map(|id| (id.web_id, id.entity_uuid)) - .collect(); - let missing_embeddings: Vec = params - .entity_ids - .into_iter() - .filter(|id| !found_set.contains(&(id.web_id, id.entity_uuid))) - .collect(); - if found_ids.is_empty() || params.cluster_count == 0 { return Ok(ClusterEntitiesResponse { clusters: Vec::new(), - missing_embeddings, + missing_embeddings: missing_ids, inertia: 0.0, }); } @@ -2831,7 +2824,7 @@ where let result = result.map_err(JoinError); - let _tx = tx.send(result); + let _: Result<(), Result> = tx.send(result); }); let clustering = rx @@ -2839,14 +2832,19 @@ where .change_context(ClusterError::Store)? .change_context(ClusterError::Store)?; - let mut groups: BTreeMap> = BTreeMap::new(); - for (index, id) in found_ids.iter().enumerate() { - groups.entry(clustering.label(index)).or_default().push(*id); + let mut groups = vec![Vec::new(); config.k as usize]; + + #[expect(clippy::indexing_slicing, reason = "we only ever have k groups")] + for (index, &id) in found_ids.iter().enumerate() { + let label = clustering.label(index) as usize; + groups[label].push(id); } let clusters = groups .into_iter() - .map(|(cluster_id, entity_ids)| EntityCluster { + .zip(0_u16..) + .filter(|(entity_ids, _)| !entity_ids.is_empty()) + .map(|(entity_ids, cluster_id)| EntityCluster { cluster_id, entity_ids, centroid: clustering.centroid(cluster_id).to_vec(), @@ -2855,7 +2853,7 @@ where Ok(ClusterEntitiesResponse { clusters, - missing_embeddings, + missing_embeddings: missing_ids, inertia: clustering.inertia, }) } diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index 063cce3a7ec..d73dd9254dc 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -581,7 +581,7 @@ pub struct ClusterEntitiesResponse { /// are omitted. pub clusters: Vec, /// Entities from the request that had no stored embedding or that do not exist. - pub missing_embeddings: Vec, + pub missing_embeddings: HashSet, /// Sum of squared chord distances from every clustered entity to its /// assigned centroid. Lower is tighter; comparable across runs over the /// same entities, e.g. to choose a cluster count. `0.0` when nothing was diff --git a/tests/graph/integration/postgres/clustering.rs b/tests/graph/integration/postgres/clustering.rs index a7016fe4229..da10b0a90cd 100644 --- a/tests/graph/integration/postgres/clustering.rs +++ b/tests/graph/integration/postgres/clustering.rs @@ -289,7 +289,7 @@ async fn permission_denied_is_reported_as_missing() { ); assert_eq!( response.missing_embeddings, - vec![entity_id], + HashSet::from([entity_id]), "unauthorized entities should be indistinguishable from missing embeddings" ); assert!(response.inertia.abs() < f32::EPSILON); From 0aa7d67377b263b3ff65ba765187010828cd4237 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:04:56 +0200 Subject: [PATCH 36/38] feat: dedupe query and make it deterministic --- libs/@local/graph/api/src/rest/mod.rs | 2 +- .../store/postgres/knowledge/entity/mod.rs | 20 ++++++++++++++----- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/libs/@local/graph/api/src/rest/mod.rs b/libs/@local/graph/api/src/rest/mod.rs index 00e6751be6d..114ac3a3748 100644 --- a/libs/@local/graph/api/src/rest/mod.rs +++ b/libs/@local/graph/api/src/rest/mod.rs @@ -596,7 +596,7 @@ where .layer(Extension(dependencies.domain_regex)) .layer(Extension(dependencies.api_config)) .layer(Extension(dependencies.compiler)) - .layer(Extension(Arc::new(dependencies.clustering))); + .layer(Extension(dependencies.clustering)); if let Some(query_logger) = dependencies.query_logger { router = router.layer(Extension(query_logger)); diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs index 4a2eca1817c..d7d388b69a8 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/mod.rs @@ -2740,14 +2740,24 @@ where .query_raw( &format!( "SELECT - e.web_id, - e.entity_uuid, + u.web_id, + u.entity_uuid, subvector(e.embedding, 1, {truncated_dim})::vector({truncated_dim}) AS \ embedding - FROM entity_embeddings e + FROM ( + SELECT DISTINCT ON (t.web_id, t.entity_uuid) + t.web_id, + t.entity_uuid, + t.ord + FROM unnest($1::uuid[], $2::uuid[]) + WITH ORDINALITY AS t(web_id, entity_uuid, ord) + ORDER BY t.web_id, t.entity_uuid, t.ord + ) u + JOIN entity_embeddings e + ON e.web_id = u.web_id + AND e.entity_uuid = u.entity_uuid WHERE e.property IS NULL - AND (e.web_id, e.entity_uuid) IN (SELECT * FROM unnest($1::uuid[], \ - $2::uuid[]))" + ORDER BY u.ord" ), [ &web_ids as &(dyn ToSql + Sync), From 90538ebe72a3ec6c343dd7df847c4f8892cd84f4 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:26:15 +0200 Subject: [PATCH 37/38] fix: docs --- libs/@local/graph/store/src/entity/store.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/@local/graph/store/src/entity/store.rs b/libs/@local/graph/store/src/entity/store.rs index d73dd9254dc..ca2a0b1458a 100644 --- a/libs/@local/graph/store/src/entity/store.rs +++ b/libs/@local/graph/store/src/entity/store.rs @@ -580,7 +580,8 @@ pub struct ClusterEntitiesResponse { /// One entry per non-empty cluster. Empty clusters (no points assigned) /// are omitted. pub clusters: Vec, - /// Entities from the request that had no stored embedding or that do not exist. + /// Entities from the request that were not clustered (either because no embedding exists, or + /// because the actor lacks permission to view the entity). pub missing_embeddings: HashSet, /// Sum of squared chord distances from every clustered entity to its /// assigned centroid. Lower is tighter; comparable across runs over the From 50be867791b929dfe03aa4a985e80f24d3f61fd2 Mon Sep 17 00:00:00 2001 From: Bilal Mahmoud <7252775+indietyp@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:34:57 +0200 Subject: [PATCH 38/38] fix: schema --- libs/@local/graph/api/openapi/openapi.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/@local/graph/api/openapi/openapi.json b/libs/@local/graph/api/openapi/openapi.json index cc55cd23940..041fe7f5621 100644 --- a/libs/@local/graph/api/openapi/openapi.json +++ b/libs/@local/graph/api/openapi/openapi.json @@ -3856,7 +3856,7 @@ "items": { "$ref": "#/components/schemas/EntityId" }, - "description": "Entities from the request that had no stored embedding or that do not exist.", + "description": "Entities from the request that were not clustered (either because no embedding exists, or\nbecause the actor lacks permission to view the entity).", "uniqueItems": true } }