1use crate::error::{SpatialError, SpatialResult};
42use crate::memory_pool::DistancePool;
43use crate::simd_distance::hardware_specific_simd::HardwareOptimizedDistances;
44use scirs2_core::ndarray::{Array1, Array2, ArrayView2};
45use scirs2_core::simd_ops::PlatformCapabilities;
46use std::collections::VecDeque;
47use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
48use std::sync::mpsc::{channel, Receiver, Sender};
49use std::sync::{Arc, Mutex};
50use std::thread;
51use std::time::Duration;
52
53#[cfg(any(target_os = "linux", target_os = "android"))]
55use libc;
56
57#[derive(Debug, Clone)]
59pub struct WorkStealingConfig {
60 pub numa_aware: bool,
62 pub work_stealing: bool,
64 pub adaptive_scheduling: bool,
66 pub num_threads: usize,
68 pub initial_chunk_size: usize,
70 pub min_chunk_size: usize,
72 pub thread_affinity: ThreadAffinityStrategy,
74 pub memory_strategy: MemoryStrategy,
76 pub prefetch_distance: usize,
78}
79
80impl Default for WorkStealingConfig {
81 fn default() -> Self {
82 Self::new()
83 }
84}
85
86impl WorkStealingConfig {
87 pub fn new() -> Self {
89 Self {
90 numa_aware: true,
91 work_stealing: true,
92 adaptive_scheduling: true,
93 num_threads: 0, initial_chunk_size: 1024,
95 min_chunk_size: 64,
96 thread_affinity: ThreadAffinityStrategy::NumaAware,
97 memory_strategy: MemoryStrategy::NumaInterleaved,
98 prefetch_distance: 8,
99 }
100 }
101
102 pub fn with_numa_aware(mut self, enabled: bool) -> Self {
104 self.numa_aware = enabled;
105 self
106 }
107
108 pub fn with_work_stealing(mut self, enabled: bool) -> Self {
110 self.work_stealing = enabled;
111 self
112 }
113
114 pub fn with_adaptive_scheduling(mut self, enabled: bool) -> Self {
116 self.adaptive_scheduling = enabled;
117 self
118 }
119
120 pub fn with_threads(mut self, numthreads: usize) -> Self {
122 self.num_threads = numthreads;
123 self
124 }
125
126 pub fn with_chunk_sizes(mut self, initial: usize, minimum: usize) -> Self {
128 self.initial_chunk_size = initial;
129 self.min_chunk_size = minimum;
130 self
131 }
132
133 pub fn with_thread_affinity(mut self, strategy: ThreadAffinityStrategy) -> Self {
135 self.thread_affinity = strategy;
136 self
137 }
138
139 pub fn with_memory_strategy(mut self, strategy: MemoryStrategy) -> Self {
141 self.memory_strategy = strategy;
142 self
143 }
144}
145
146#[derive(Debug, Clone, PartialEq)]
148pub enum ThreadAffinityStrategy {
149 None,
151 Physical,
153 NumaAware,
155 Custom(Vec<usize>),
157}
158
159#[derive(Debug, Clone, PartialEq)]
161pub enum MemoryStrategy {
162 System,
164 NumaLocal,
166 NumaInterleaved,
168 HugePages,
170}
171
172#[derive(Debug, Clone)]
174pub struct NumaTopology {
175 pub num_nodes: usize,
177 pub cores_per_node: Vec<usize>,
179 pub memory_per_node: Vec<usize>,
181 pub distance_matrix: Vec<Vec<u32>>,
183}
184
185impl Default for NumaTopology {
186 fn default() -> Self {
187 Self::detect()
188 }
189}
190
191impl NumaTopology {
192 pub fn detect() -> Self {
206 #[cfg(target_os = "linux")]
207 {
208 if let Some(topology) = Self::detect_linux() {
209 return topology;
210 }
211 }
212
213 Self::single_node_fallback()
214 }
215
216 fn single_node_fallback() -> Self {
221 let logical_cpus = thread::available_parallelism()
222 .map(|n| n.get())
223 .unwrap_or(1)
224 .max(1);
225
226 Self {
227 num_nodes: 1,
228 cores_per_node: vec![logical_cpus],
229 memory_per_node: vec![0], distance_matrix: Self::create_default_distance_matrix(1),
231 }
232 }
233
234 #[cfg(target_os = "linux")]
240 fn detect_linux() -> Option<Self> {
241 let node_root = std::path::Path::new("/sys/devices/system/node");
242 let entries = std::fs::read_dir(node_root).ok()?;
243
244 let mut node_ids: Vec<usize> = entries
246 .filter_map(|entry| entry.ok())
247 .filter_map(|entry| {
248 let name = entry.file_name();
249 let name = name.to_str()?;
250 name.strip_prefix("node")
251 .and_then(|suffix| suffix.parse::<usize>().ok())
252 })
253 .collect();
254 node_ids.sort_unstable();
255
256 if node_ids.is_empty() {
257 return None;
258 }
259
260 let num_nodes = node_ids.len();
261 let mut cores_per_node = Vec::with_capacity(num_nodes);
262 let mut memory_per_node = Vec::with_capacity(num_nodes);
263
264 for &node_id in &node_ids {
265 let cpulist_path = node_root.join(format!("node{node_id}/cpulist"));
270 let core_count = std::fs::read_to_string(&cpulist_path)
271 .ok()
272 .map(|list| Self::count_cpus_in_list(list.trim()))
273 .unwrap_or(0)
274 .max(1);
275 cores_per_node.push(core_count);
276
277 let meminfo_path = node_root.join(format!("node{node_id}/meminfo"));
280 let memory_bytes = std::fs::read_to_string(&meminfo_path)
281 .ok()
282 .and_then(|content| {
283 content.lines().find_map(|line| {
284 if line.contains("MemTotal:") {
285 line.split_whitespace()
286 .rev()
287 .nth(1)
288 .and_then(|kib| kib.parse::<usize>().ok())
289 .map(|kib| kib.saturating_mul(1024))
290 } else {
291 None
292 }
293 })
294 })
295 .unwrap_or(0); memory_per_node.push(memory_bytes);
297 }
298
299 let distance_matrix = Self::read_distance_matrix(node_root, &node_ids)
300 .unwrap_or_else(|| Self::create_default_distance_matrix(num_nodes));
301
302 Some(Self {
303 num_nodes,
304 cores_per_node,
305 memory_per_node,
306 distance_matrix,
307 })
308 }
309
310 #[cfg(target_os = "linux")]
318 fn read_distance_matrix(
319 node_root: &std::path::Path,
320 node_ids: &[usize],
321 ) -> Option<Vec<Vec<u32>>> {
322 let num_nodes = node_ids.len();
323 let mut matrix = Vec::with_capacity(num_nodes);
324
325 for &node_id in node_ids {
326 let distance_path = node_root.join(format!("node{node_id}/distance"));
327 let contents = std::fs::read_to_string(&distance_path).ok()?;
328 let row: Vec<u32> = contents
329 .split_whitespace()
330 .map(|value| value.parse::<u32>().ok())
331 .collect::<Option<Vec<u32>>>()?;
332
333 if row.len() != num_nodes {
334 return None;
335 }
336 matrix.push(row);
337 }
338
339 Some(matrix)
340 }
341
342 #[cfg(target_os = "linux")]
345 fn count_cpus_in_list(list: &str) -> usize {
346 if list.is_empty() {
347 return 0;
348 }
349
350 let mut count = 0usize;
351 for part in list.split(',') {
352 if let Some((start, end)) = part.split_once('-') {
353 if let (Ok(start), Ok(end)) = (start.parse::<usize>(), end.parse::<usize>()) {
354 if end >= start {
355 count += end - start + 1;
356 }
357 }
358 } else if part.parse::<usize>().is_ok() {
359 count += 1;
360 }
361 }
362 count
363 }
364
365 #[allow(clippy::needless_range_loop)]
369 fn create_default_distance_matrix(num_nodes: usize) -> Vec<Vec<u32>> {
370 let mut matrix = vec![vec![0; num_nodes]; num_nodes];
371 for i in 0..num_nodes {
372 for j in 0..num_nodes {
373 if i == j {
374 matrix[i][j] = 10; } else {
376 matrix[i][j] = 20; }
378 }
379 }
380 matrix
381 }
382
383 pub fn optimal_threads_per_node(&self, node: usize) -> usize {
385 if node < self.cores_per_node.len() {
386 self.cores_per_node[node]
387 } else {
388 self.cores_per_node.first().copied().unwrap_or(1)
389 }
390 }
391
392 pub fn memory_capacity(&self, node: usize) -> usize {
394 self.memory_per_node.get(node).copied().unwrap_or(0)
395 }
396}
397
398pub struct WorkStealingPool {
400 workers: Vec<WorkStealingWorker>,
401 #[allow(dead_code)]
402 config: WorkStealingConfig,
403 numa_topology: NumaTopology,
404 global_queue: Arc<Mutex<VecDeque<WorkItem>>>,
405 completed_work: Arc<AtomicUsize>,
406 total_work: Arc<AtomicUsize>,
407 active_workers: Arc<AtomicUsize>,
408 shutdown: Arc<AtomicBool>,
409}
410
411struct WorkStealingWorker {
413 thread_id: usize,
414 numa_node: usize,
415 local_queue: Arc<Mutex<VecDeque<WorkItem>>>,
416 thread_handle: Option<thread::JoinHandle<()>>,
417 memory_pool: Arc<DistancePool>,
418}
419
420#[derive(Debug, Clone)]
422pub struct WorkItem {
423 pub start: usize,
425 pub end: usize,
427 pub work_type: WorkType,
429 pub priority: u8,
431 pub numa_hint: Option<usize>,
433}
434
435#[derive(Debug, Clone, PartialEq)]
437pub enum WorkType {
438 DistanceMatrix,
440 KMeansClustering,
442 KDTreeBuild,
444 NearestNeighbor,
446 Custom(String),
448}
449
450pub struct WorkContext {
452 pub distance_context: Option<DistanceMatrixContext>,
454 pub kmeans_context: Option<KMeansContext>,
456 pub kdtree_context: Option<KDTreeContext>,
458 pub nn_context: Option<NearestNeighborContext>,
460 pub custom_context: Option<CustomWorkContext>,
462}
463
464pub struct DistanceMatrixContext {
466 pub points: Array2<f64>,
468 pub result_sender: Sender<(usize, usize, f64)>,
470}
471
472pub struct KMeansContext {
474 pub points: Array2<f64>,
476 pub centroids: Array2<f64>,
478 pub assignment_sender: Sender<(usize, usize)>,
480}
481
482pub struct KDTreeContext {
484 pub points: Array2<f64>,
486 pub indices: Vec<usize>,
488 pub depth: usize,
490 pub config: KDTreeConfig,
492 pub result_sender: Sender<(usize, KDTreeChunkResult)>,
494}
495
496pub struct NearestNeighborContext {
498 pub query_points: Array2<f64>,
500 pub data_points: Array2<f64>,
502 pub k: usize,
504 pub result_sender: Sender<(usize, Vec<(usize, f64)>)>,
506}
507
508pub struct CustomWorkContext {
510 pub process_fn: fn(usize, usize, &CustomUserData),
512 pub user_data: CustomUserData,
514}
515
516#[derive(Debug, Clone)]
518pub struct CustomUserData {
519 pub data: Vec<u8>,
521}
522
523#[derive(Debug, Clone)]
525pub struct KDTreeConfig {
526 pub max_leaf_size: usize,
528 pub cache_aware: bool,
530}
531
532impl Default for KDTreeConfig {
533 fn default() -> Self {
534 Self {
535 max_leaf_size: 32,
536 cache_aware: true,
537 }
538 }
539}
540
541#[derive(Debug, Clone)]
543pub struct KDTreeChunkResult {
544 pub node_index: usize,
546 pub is_leaf: bool,
548 pub splitting_dimension: usize,
550 pub split_value: f64,
552 pub left_indices: Vec<usize>,
554 pub right_indices: Vec<usize>,
556}
557
558impl WorkStealingPool {
559 pub fn new(config: WorkStealingConfig) -> SpatialResult<Self> {
561 let numa_topology = if config.numa_aware {
562 NumaTopology::detect()
563 } else {
564 NumaTopology {
565 num_nodes: 1,
566 cores_per_node: vec![config.num_threads],
567 memory_per_node: vec![0],
568 distance_matrix: vec![vec![10]],
569 }
570 };
571
572 let num_threads = if config.num_threads == 0 {
573 numa_topology.cores_per_node.iter().sum()
574 } else {
575 config.num_threads
576 };
577
578 let global_queue = Arc::new(Mutex::new(VecDeque::new()));
579 let completed_work = Arc::new(AtomicUsize::new(0));
580 let total_work = Arc::new(AtomicUsize::new(0));
581 let active_workers = Arc::new(AtomicUsize::new(0));
582 let shutdown = Arc::new(AtomicBool::new(false));
583
584 let mut workers = Vec::with_capacity(num_threads);
585
586 for thread_id in 0..num_threads {
588 let numa_node = if config.numa_aware {
589 Self::assign_thread_to_numa_node(thread_id, &numa_topology)
590 } else {
591 0
592 };
593
594 let worker = WorkStealingWorker {
595 thread_id,
596 numa_node,
597 local_queue: Arc::new(Mutex::new(VecDeque::new())),
598 thread_handle: None,
599 memory_pool: Arc::new(DistancePool::new(1000)),
600 };
601
602 workers.push(worker);
603 }
604
605 for worker in &mut workers {
607 let local_queue = Arc::clone(&worker.local_queue);
608 let global_queue = Arc::clone(&global_queue);
609 let completed_work = Arc::clone(&completed_work);
610 let active_workers = Arc::clone(&active_workers);
611 let shutdown = Arc::clone(&shutdown);
612 let config_clone = config.clone();
613 let thread_id = worker.thread_id;
614 let numa_node = worker.numa_node;
615 let memory_pool = Arc::clone(&worker.memory_pool);
616
617 let handle = thread::spawn(move || {
618 Self::worker_main(
619 thread_id,
620 numa_node,
621 local_queue,
622 global_queue,
623 completed_work,
624 active_workers,
625 shutdown,
626 config_clone,
627 memory_pool,
628 );
629 });
630
631 worker.thread_handle = Some(handle);
632 }
633
634 Ok(Self {
635 workers,
636 config,
637 numa_topology,
638 global_queue,
639 completed_work,
640 total_work,
641 active_workers,
642 shutdown,
643 })
644 }
645
646 fn assign_thread_to_numa_node(_threadid: usize, topology: &NumaTopology) -> usize {
648 let mut thread_count = 0;
649 for (node_id, &cores) in topology.cores_per_node.iter().enumerate() {
650 if _threadid < thread_count + cores {
651 return node_id;
652 }
653 thread_count += cores;
654 }
655 0 }
657
658 fn worker_main(
660 thread_id: usize,
661 numa_node: usize,
662 local_queue: Arc<Mutex<VecDeque<WorkItem>>>,
663 global_queue: Arc<Mutex<VecDeque<WorkItem>>>,
664 completed_work: Arc<AtomicUsize>,
665 active_workers: Arc<AtomicUsize>,
666 shutdown: Arc<AtomicBool>,
667 config: WorkStealingConfig,
668 memory_pool: Arc<DistancePool>,
669 ) {
670 Self::set_thread_affinity(thread_id, numa_node, &config);
672
673 let work_context = WorkContext {
675 distance_context: None,
676 kmeans_context: None,
677 kdtree_context: None,
678 nn_context: None,
679 custom_context: None,
680 };
681
682 while !shutdown.load(Ordering::Relaxed) {
683 let work_item = Self::get_work_item(&local_queue, &global_queue, &config);
684
685 if let Some(item) = work_item {
686 active_workers.fetch_add(1, Ordering::Relaxed);
687
688 Self::process_work_item(item, &work_context);
690
691 completed_work.fetch_add(1, Ordering::Relaxed);
692 active_workers.fetch_sub(1, Ordering::Relaxed);
693 } else {
694 if config.work_stealing {
696 Self::attempt_work_stealing(thread_id, &local_queue, &global_queue, &config);
697 }
698
699 thread::sleep(Duration::from_micros(100));
701 }
702 }
703 }
704
705 fn set_thread_affinity(thread_id: usize, numanode: usize, config: &WorkStealingConfig) {
707 match config.thread_affinity {
708 ThreadAffinityStrategy::Physical => {
709 #[cfg(target_os = "linux")]
712 {
713 if let Err(e) = Self::set_cpu_affinity_linux(thread_id) {
714 eprintln!(
715 "Warning: Failed to set CPU affinity for thread {thread_id}: {e}"
716 );
717 }
718 }
719 #[cfg(target_os = "windows")]
720 {
721 if let Err(e) = Self::set_cpu_affinity_windows(thread_id) {
722 eprintln!(
723 "Warning: Failed to set CPU affinity for thread {}: {}",
724 thread_id, e
725 );
726 }
727 }
728 }
729 ThreadAffinityStrategy::NumaAware => {
730 #[cfg(target_os = "linux")]
732 {
733 if let Err(e) = Self::set_numa_affinity_linux(numanode) {
734 eprintln!(
735 "Warning: Failed to set NUMA affinity for node {}: {}",
736 numanode, e
737 );
738 }
739 }
740 #[cfg(target_os = "windows")]
741 {
742 if let Err(e) = Self::set_numa_affinity_windows(numanode) {
743 eprintln!(
744 "Warning: Failed to set NUMA affinity for node {}: {}",
745 numanode, e
746 );
747 }
748 }
749 }
750 ThreadAffinityStrategy::Custom(ref cpus) => {
751 if let Some(&cpu) = cpus.get(thread_id) {
752 #[cfg(target_os = "linux")]
753 {
754 if let Err(e) = Self::set_custom_cpu_affinity_linux(cpu) {
755 eprintln!(
756 "Warning: Failed to set custom CPU affinity to core {cpu}: {e}"
757 );
758 }
759 }
760 #[cfg(target_os = "windows")]
761 {
762 if let Err(e) = Self::set_custom_cpu_affinity_windows(cpu) {
763 eprintln!(
764 "Warning: Failed to set custom CPU affinity to core {}: {}",
765 cpu, e
766 );
767 }
768 }
769 }
770 }
771 ThreadAffinityStrategy::None => {
772 }
774 }
775 }
776
777 #[cfg(target_os = "linux")]
779 fn set_cpu_affinity_linux(_cpuid: usize) -> Result<(), Box<dyn std::error::Error>> {
780 unsafe {
781 let mut cpu_set: libc::cpu_set_t = std::mem::zeroed();
782 libc::CPU_SET(_cpuid, &mut cpu_set);
783
784 let result = libc::sched_setaffinity(
785 0, std::mem::size_of::<libc::cpu_set_t>(),
787 &cpu_set,
788 );
789
790 if result == 0 {
791 Ok(())
792 } else {
793 Err("Failed to set CPU affinity".into())
794 }
795 }
796 }
797
798 #[cfg(target_os = "linux")]
800 fn set_numa_affinity_linux(_numanode: usize) -> Result<(), Box<dyn std::error::Error>> {
801 use std::fs;
802
803 let cpulist_path = format!("/sys/devices/system/node/node{}/cpulist", _numanode);
805 let cpulist = fs::read_to_string(&cpulist_path)
806 .map_err(|_| format!("Failed to read NUMA node {} CPU list", _numanode))?;
807
808 unsafe {
809 let mut cpu_set: libc::cpu_set_t = std::mem::zeroed();
810
811 for range in cpulist.trim().split(',') {
813 if let Some((start, end)) = range.split_once('-') {
814 if let (Ok(s), Ok(e)) = (start.parse::<u32>(), end.parse::<u32>()) {
815 for cpu in s..=e {
816 libc::CPU_SET(cpu as usize, &mut cpu_set);
817 }
818 }
819 } else if let Ok(cpu) = range.parse::<u32>() {
820 libc::CPU_SET(cpu as usize, &mut cpu_set);
821 }
822 }
823
824 let result = libc::sched_setaffinity(
825 0, std::mem::size_of::<libc::cpu_set_t>(),
827 &cpu_set,
828 );
829
830 if result == 0 {
831 Ok(())
832 } else {
833 Err("Failed to set NUMA affinity".into())
834 }
835 }
836 }
837
838 #[cfg(target_os = "linux")]
840 fn set_custom_cpu_affinity_linux(_cpuid: usize) -> Result<(), Box<dyn std::error::Error>> {
841 Self::set_cpu_affinity_linux(_cpuid)
843 }
844
845 #[cfg(target_os = "windows")]
847 fn set_cpu_affinity_windows(_cpuid: usize) -> Result<(), Box<dyn std::error::Error>> {
848 let _ = _cpuid;
851 Ok(())
852 }
853
854 #[cfg(target_os = "windows")]
856 fn set_numa_affinity_windows(_numanode: usize) -> Result<(), Box<dyn std::error::Error>> {
857 let _ = _numanode;
860 Ok(())
861 }
862
863 #[cfg(target_os = "windows")]
865 fn set_custom_cpu_affinity_windows(_cpuid: usize) -> Result<(), Box<dyn std::error::Error>> {
866 Self::set_cpu_affinity_windows(_cpuid)
868 }
869
870 fn get_work_item(
872 local_queue: &Arc<Mutex<VecDeque<WorkItem>>>,
873 global_queue: &Arc<Mutex<VecDeque<WorkItem>>>,
874 config: &WorkStealingConfig,
875 ) -> Option<WorkItem> {
876 if let Ok(mut queue) = local_queue.try_lock() {
878 if let Some(item) = queue.pop_front() {
879 return Some(item);
880 }
881 }
882
883 if let Ok(mut queue) = global_queue.try_lock() {
885 if let Some(item) = queue.pop_front() {
886 return Some(item);
887 }
888 }
889
890 None
891 }
892
893 fn attempt_work_stealing(
895 _threadid: usize,
896 _queue: &Arc<Mutex<VecDeque<WorkItem>>>,
897 _global_queue: &Arc<Mutex<VecDeque<WorkItem>>>,
898 config: &WorkStealingConfig,
899 ) {
900 }
903
904 fn process_work_item(item: WorkItem, context: &WorkContext) {
906 match item.work_type {
907 WorkType::DistanceMatrix => {
908 Self::process_distance_matrix_chunk(item.start, item.end, context);
909 }
910 WorkType::KMeansClustering => {
911 Self::process_kmeans_chunk(item.start, item.end, context);
912 }
913 WorkType::KDTreeBuild => {
914 Self::process_kdtree_chunk(item.start, item.end, context);
915 }
916 WorkType::NearestNeighbor => {
917 Self::process_nn_chunk(item.start, item.end, context);
918 }
919 WorkType::Custom(_name) => {
920 Self::process_custom_chunk(item.start, item.end, context);
921 }
922 }
923 }
924
925 fn process_distance_matrix_chunk(start: usize, end: usize, context: &WorkContext) {
927 if let Some(distance_context) = &context.distance_context {
928 use crate::simd_distance::hardware_specific_simd::HardwareOptimizedDistances;
929
930 let optimizer = HardwareOptimizedDistances::new();
931 let points = &distance_context.points;
932 let n_points = points.nrows();
933
934 for _linearidx in start..end {
936 let (i, j) = Self::linear_to_matrix_indices(_linearidx, n_points);
937
938 if i < j && i < n_points && j < n_points {
939 let point_i = points.row(i);
940 let point_j = points.row(j);
941
942 match optimizer.euclidean_distance_optimized(&point_i, &point_j) {
943 Ok(distance) => {
944 distance_context.result_sender.send((i, j, distance)).ok();
946 }
947 Err(_) => {
948 distance_context.result_sender.send((i, j, f64::NAN)).ok();
950 }
951 }
952 }
953 }
954 }
955 }
956
957 fn process_kmeans_chunk(start: usize, end: usize, context: &WorkContext) {
959 if let Some(kmeans_context) = &context.kmeans_context {
960 let optimizer = HardwareOptimizedDistances::new();
961 let points = &kmeans_context.points;
962 let centroids = &kmeans_context.centroids;
963 let k = centroids.nrows();
964
965 for point_idx in start..end {
967 if point_idx < points.nrows() {
968 let point = points.row(point_idx);
969 let mut best_cluster = 0;
970 let mut best_distance = f64::INFINITY;
971
972 for cluster_idx in 0..k {
974 let centroid = centroids.row(cluster_idx);
975
976 match optimizer.euclidean_distance_optimized(&point, ¢roid) {
977 Ok(distance) => {
978 if distance < best_distance {
979 best_distance = distance;
980 best_cluster = cluster_idx;
981 }
982 }
983 Err(_) => continue,
984 }
985 }
986
987 kmeans_context
989 .assignment_sender
990 .send((point_idx, best_cluster))
991 .ok();
992 }
993 }
994 }
995 }
996
997 fn process_kdtree_chunk(start: usize, end: usize, context: &WorkContext) {
999 if let Some(kdtree_context) = &context.kdtree_context {
1000 let points = &kdtree_context.points;
1001 let indices = &kdtree_context.indices;
1002 let depth = kdtree_context.depth;
1003
1004 let chunk_indices: Vec<usize> = indices[start..end.min(indices.len())].to_vec();
1006
1007 if !chunk_indices.is_empty() {
1008 let local_tree = Self::build_local_kdtree_chunk(
1010 points,
1011 &chunk_indices,
1012 depth,
1013 &kdtree_context.config,
1014 );
1015
1016 kdtree_context.result_sender.send((start, local_tree)).ok();
1018 }
1019 }
1020 }
1021
1022 fn process_nn_chunk(start: usize, end: usize, context: &WorkContext) {
1024 if let Some(nn_context) = &context.nn_context {
1025 let optimizer = HardwareOptimizedDistances::new();
1026 let query_points = &nn_context.query_points;
1027 let data_points = &nn_context.data_points;
1028 let k = nn_context.k;
1029
1030 for query_idx in start..end {
1032 if query_idx < query_points.nrows() {
1033 let query = query_points.row(query_idx);
1034
1035 let mut distances: Vec<(f64, usize)> = Vec::with_capacity(data_points.nrows());
1037
1038 for (data_idx, data_point) in data_points.outer_iter().enumerate() {
1039 match optimizer.euclidean_distance_optimized(&query, &data_point) {
1040 Ok(distance) => distances.push((distance, data_idx)),
1041 Err(_) => distances.push((f64::INFINITY, data_idx)),
1042 }
1043 }
1044
1045 if k <= distances.len() {
1047 distances.select_nth_unstable_by(k - 1, |a, b| {
1048 a.0.partial_cmp(&b.0).expect("Operation failed")
1049 });
1050 distances[..k].sort_unstable_by(|a, b| {
1051 a.0.partial_cmp(&b.0).expect("Operation failed")
1052 });
1053
1054 let result: Vec<(usize, f64)> = distances[..k]
1055 .iter()
1056 .map(|(dist, idx)| (*idx, *dist))
1057 .collect();
1058
1059 nn_context.result_sender.send((query_idx, result)).ok();
1060 }
1061 }
1062 }
1063 }
1064 }
1065
1066 fn process_custom_chunk(start: usize, end: usize, context: &WorkContext) {
1068 if let Some(custom_context) = &context.custom_context {
1069 (custom_context.process_fn)(start, end, &custom_context.user_data);
1071 }
1072 }
1073
1074 fn linear_to_matrix_indices(_linearidx: usize, n: usize) -> (usize, usize) {
1076 let mut k = _linearidx;
1078 let mut i = 0;
1079
1080 while k >= n - i - 1 {
1081 k -= n - i - 1;
1082 i += 1;
1083 }
1084
1085 let j = k + i + 1;
1086 (i, j)
1087 }
1088
1089 fn build_local_kdtree_chunk(
1091 points: &Array2<f64>,
1092 indices: &[usize],
1093 depth: usize,
1094 config: &KDTreeConfig,
1095 ) -> KDTreeChunkResult {
1096 let n_dims = points.ncols();
1097 let splitting_dimension = depth % n_dims;
1098
1099 if indices.len() <= 1 {
1100 return KDTreeChunkResult {
1101 node_index: indices.first().copied().unwrap_or(0),
1102 is_leaf: true,
1103 splitting_dimension,
1104 split_value: 0.0,
1105 left_indices: Vec::new(),
1106 right_indices: Vec::new(),
1107 };
1108 }
1109
1110 let mut sorted_indices = indices.to_vec();
1112 sorted_indices.sort_by(|&a, &b| {
1113 let coord_a = points[[a, splitting_dimension]];
1114 let coord_b = points[[b, splitting_dimension]];
1115 coord_a
1116 .partial_cmp(&coord_b)
1117 .unwrap_or(std::cmp::Ordering::Equal)
1118 });
1119
1120 let median_idx = sorted_indices.len() / 2;
1121 let split_point_idx = sorted_indices[median_idx];
1122 let split_value = points[[split_point_idx, splitting_dimension]];
1123
1124 let left_indices = sorted_indices[..median_idx].to_vec();
1125 let right_indices = sorted_indices[median_idx + 1..].to_vec();
1126
1127 KDTreeChunkResult {
1128 node_index: split_point_idx,
1129 is_leaf: false,
1130 splitting_dimension,
1131 split_value,
1132 left_indices,
1133 right_indices,
1134 }
1135 }
1136
1137 pub fn submit_work(&self, _workitems: Vec<WorkItem>) -> SpatialResult<()> {
1139 self.total_work.store(_workitems.len(), Ordering::Relaxed);
1140 self.completed_work.store(0, Ordering::Relaxed);
1141
1142 let mut global_queue = self.global_queue.lock().expect("Operation failed");
1143 for item in _workitems {
1144 global_queue.push_back(item);
1145 }
1146 drop(global_queue);
1147
1148 Ok(())
1149 }
1150
1151 pub fn wait_for_completion(&self) -> SpatialResult<()> {
1153 let total = self.total_work.load(Ordering::Relaxed);
1154
1155 while self.completed_work.load(Ordering::Relaxed) < total {
1156 thread::sleep(Duration::from_millis(1));
1157 }
1158
1159 Ok(())
1160 }
1161
1162 pub fn progress(&self) -> (usize, usize) {
1164 let completed = self.completed_work.load(Ordering::Relaxed);
1165 let total = self.total_work.load(Ordering::Relaxed);
1166 (completed, total)
1167 }
1168
1169 pub fn statistics(&self) -> PoolStatistics {
1171 PoolStatistics {
1172 num_threads: self.workers.len(),
1173 numa_nodes: self.numa_topology.num_nodes,
1174 active_workers: self.active_workers.load(Ordering::Relaxed),
1175 completed_work: self.completed_work.load(Ordering::Relaxed),
1176 total_work: self.total_work.load(Ordering::Relaxed),
1177 queue_depth: self.global_queue.lock().expect("Operation failed").len(),
1178 }
1179 }
1180}
1181
1182impl Drop for WorkStealingPool {
1183 fn drop(&mut self) {
1184 self.shutdown.store(true, Ordering::Relaxed);
1186
1187 for worker in &mut self.workers {
1189 if let Some(handle) = worker.thread_handle.take() {
1190 let _ = handle.join();
1191 }
1192 }
1193 }
1194}
1195
1196#[derive(Debug, Clone)]
1198pub struct PoolStatistics {
1199 pub num_threads: usize,
1200 pub numa_nodes: usize,
1201 pub active_workers: usize,
1202 pub completed_work: usize,
1203 pub total_work: usize,
1204 pub queue_depth: usize,
1205}
1206
1207pub struct AdvancedParallelDistanceMatrix {
1209 pool: WorkStealingPool,
1210 config: WorkStealingConfig,
1211}
1212
1213impl AdvancedParallelDistanceMatrix {
1214 pub fn new(config: WorkStealingConfig) -> SpatialResult<Self> {
1216 let pool = WorkStealingPool::new(config.clone())?;
1217 Ok(Self { pool, config })
1218 }
1219
1220 pub fn compute_parallel(&self, points: &ArrayView2<'_, f64>) -> SpatialResult<Array2<f64>> {
1222 let n_points = points.nrows();
1223 let n_pairs = n_points * (n_points - 1) / 2;
1224 let mut result_matrix = Array2::zeros((n_points, n_points));
1225
1226 type DistanceResult = (usize, usize, f64);
1228 let (result_sender, result_receiver): (Sender<DistanceResult>, Receiver<DistanceResult>) =
1229 channel();
1230
1231 let _distance_context = DistanceMatrixContext {
1233 points: points.to_owned(),
1234 result_sender,
1235 };
1236
1237 let chunk_size = self.config.initial_chunk_size;
1242 let mut work_items = Vec::new();
1243
1244 for chunk_start in (0..n_pairs).step_by(chunk_size) {
1245 let chunk_end = (chunk_start + chunk_size).min(n_pairs);
1246 work_items.push(WorkItem {
1247 start: chunk_start,
1248 end: chunk_end,
1249 work_type: WorkType::DistanceMatrix,
1250 priority: 1,
1251 numa_hint: None,
1252 });
1253 }
1254
1255 self.pool.submit_work(work_items)?;
1257
1258 let mut collected_results = 0;
1260 let timeout = Duration::from_secs(2); let start_time = std::time::Instant::now();
1262
1263 while collected_results < n_pairs && start_time.elapsed() < timeout {
1264 if let Ok((i, j, distance)) = result_receiver.try_recv() {
1265 if i < n_points && j < n_points {
1266 result_matrix[[i, j]] = distance;
1267 result_matrix[[j, i]] = distance;
1268 collected_results += 1;
1269 }
1270 } else {
1271 thread::sleep(Duration::from_millis(1));
1272 }
1273 }
1274
1275 self.pool.wait_for_completion()?;
1277
1278 if collected_results < n_pairs {
1280 let optimizer = HardwareOptimizedDistances::new();
1281
1282 for i in 0..n_points {
1283 for j in (i + 1)..n_points {
1284 if result_matrix[[i, j]] == 0.0 && i != j {
1285 let point_i = points.row(i);
1286 let point_j = points.row(j);
1287
1288 if let Ok(distance) =
1289 optimizer.euclidean_distance_optimized(&point_i, &point_j)
1290 {
1291 result_matrix[[i, j]] = distance;
1292 result_matrix[[j, i]] = distance;
1293 }
1294 }
1295 }
1296 }
1297 }
1298
1299 Ok(result_matrix)
1300 }
1301
1302 pub fn statistics(&self) -> PoolStatistics {
1304 self.pool.statistics()
1305 }
1306}
1307
1308pub struct AdvancedParallelKMeans {
1310 pool: WorkStealingPool,
1311 config: WorkStealingConfig,
1312 k: usize,
1313}
1314
1315impl AdvancedParallelKMeans {
1316 pub fn new(k: usize, config: WorkStealingConfig) -> SpatialResult<Self> {
1318 let pool = WorkStealingPool::new(config.clone())?;
1319 Ok(Self { pool, config, k })
1320 }
1321
1322 pub fn fit_parallel(
1334 &self,
1335 points: &ArrayView2<'_, f64>,
1336 ) -> SpatialResult<(Array2<f64>, Array1<usize>)> {
1337 let n_points = points.nrows();
1338 let n_dims = points.ncols();
1339
1340 if self.k == 0 {
1341 return Err(SpatialError::ValueError(
1342 "Number of clusters k must be greater than zero".to_string(),
1343 ));
1344 }
1345 if n_points == 0 {
1346 return Err(SpatialError::ValueError(
1347 "Cannot cluster an empty point set".to_string(),
1348 ));
1349 }
1350 if self.k > n_points {
1351 return Err(SpatialError::ValueError(format!(
1352 "Number of clusters k ({}) cannot exceed number of points ({})",
1353 self.k, n_points
1354 )));
1355 }
1356
1357 let optimizer = HardwareOptimizedDistances::new();
1358
1359 let mut centroids = Array2::<f64>::zeros((self.k, n_dims));
1364 centroids.row_mut(0).assign(&points.row(0));
1365
1366 let squared_distance = |a: &scirs2_core::ndarray::ArrayView1<f64>,
1367 b: &scirs2_core::ndarray::ArrayView1<f64>|
1368 -> f64 {
1369 match optimizer.euclidean_distance_optimized(a, b) {
1370 Ok(d) => d * d,
1371 Err(_) => f64::INFINITY,
1372 }
1373 };
1374
1375 let mut nearest_sq = vec![f64::INFINITY; n_points];
1376 for chosen in 1..self.k {
1377 let prev_centroid = centroids.row(chosen - 1);
1378 let mut best_point = 0usize;
1379 let mut best_dist = -1.0f64;
1380 for point_idx in 0..n_points {
1381 let point = points.row(point_idx);
1382 let d_sq = squared_distance(&point, &prev_centroid);
1383 if d_sq < nearest_sq[point_idx] {
1384 nearest_sq[point_idx] = d_sq;
1385 }
1386 if nearest_sq[point_idx] > best_dist {
1387 best_dist = nearest_sq[point_idx];
1388 best_point = point_idx;
1389 }
1390 }
1391 centroids.row_mut(chosen).assign(&points.row(best_point));
1392 }
1393
1394 let mut assignments = Array1::<usize>::zeros(n_points);
1395 let max_iterations = 100;
1396
1397 for _iteration in 0..max_iterations {
1398 let mut changed = false;
1400 for point_idx in 0..n_points {
1401 let point = points.row(point_idx);
1402 let mut best_cluster = 0usize;
1403 let mut best_distance = f64::INFINITY;
1404 for cluster_idx in 0..self.k {
1405 let centroid = centroids.row(cluster_idx);
1406 let d_sq = squared_distance(&point, ¢roid);
1407 if d_sq < best_distance {
1408 best_distance = d_sq;
1409 best_cluster = cluster_idx;
1410 }
1411 }
1412 if assignments[point_idx] != best_cluster {
1413 assignments[point_idx] = best_cluster;
1414 changed = true;
1415 }
1416 }
1417
1418 let mut new_centroids = Array2::<f64>::zeros((self.k, n_dims));
1420 let mut counts = vec![0usize; self.k];
1421 for point_idx in 0..n_points {
1422 let cluster = assignments[point_idx];
1423 counts[cluster] += 1;
1424 let point = points.row(point_idx);
1425 for dim in 0..n_dims {
1426 new_centroids[[cluster, dim]] += point[dim];
1427 }
1428 }
1429 for cluster_idx in 0..self.k {
1430 if counts[cluster_idx] > 0 {
1431 let inv = 1.0 / counts[cluster_idx] as f64;
1432 for dim in 0..n_dims {
1433 new_centroids[[cluster_idx, dim]] *= inv;
1434 }
1435 } else {
1436 new_centroids
1439 .row_mut(cluster_idx)
1440 .assign(¢roids.row(cluster_idx));
1441 }
1442 }
1443 centroids = new_centroids;
1444
1445 if !changed {
1447 break;
1448 }
1449 }
1450
1451 Ok((centroids, assignments))
1452 }
1453}
1454
1455static GLOBAL_WORK_STEALING_POOL: std::sync::OnceLock<Mutex<Option<WorkStealingPool>>> =
1457 std::sync::OnceLock::new();
1458
1459#[allow(dead_code)]
1461pub fn global_work_stealing_pool() -> SpatialResult<&'static Mutex<Option<WorkStealingPool>>> {
1462 Ok(GLOBAL_WORK_STEALING_POOL.get_or_init(|| Mutex::new(None)))
1463}
1464
1465#[allow(dead_code)]
1467pub fn initialize_global_pool(config: WorkStealingConfig) -> SpatialResult<()> {
1468 let pool_mutex = global_work_stealing_pool()?;
1469 let mut pool_guard = pool_mutex.lock().expect("Operation failed");
1470
1471 if pool_guard.is_none() {
1472 *pool_guard = Some(WorkStealingPool::new(config)?);
1473 }
1474
1475 Ok(())
1476}
1477
1478#[allow(dead_code)]
1480pub fn get_numa_topology() -> NumaTopology {
1481 NumaTopology::detect()
1482}
1483
1484#[allow(dead_code)]
1486pub fn report_advanced_parallel_capabilities() {
1487 let topology = get_numa_topology();
1488 let total_cores: usize = topology.cores_per_node.iter().sum();
1489
1490 println!("Advanced-Parallel Processing Capabilities:");
1491 println!(" Total CPU cores: {total_cores}");
1492 println!(" NUMA nodes: {}", topology.num_nodes);
1493
1494 for (node, &cores) in topology.cores_per_node.iter().enumerate() {
1495 let memory_gb = topology.memory_per_node[node] as f64 / (1024.0 * 1024.0 * 1024.0);
1496 println!(" Node {node}: {cores} cores, {memory_gb:.1} GB memory");
1497 }
1498
1499 println!(" Work-stealing: Available");
1500 println!(" NUMA-aware allocation: Available");
1501 println!(" Thread affinity: Available");
1502
1503 let caps = PlatformCapabilities::detect();
1504 if caps.simd_available {
1505 println!(" SIMD acceleration: Available");
1506 if caps.avx512_available {
1507 println!(" AVX-512: Available");
1508 } else if caps.avx2_available {
1509 println!(" AVX2: Available");
1510 }
1511 }
1512}
1513
1514#[cfg(test)]
1515mod tests {
1516 use super::*;
1517 use scirs2_core::ndarray::array;
1518
1519 #[test]
1520 fn test_work_stealing_config() {
1521 let config = WorkStealingConfig::new()
1522 .with_numa_aware(true)
1523 .with_work_stealing(true)
1524 .with_threads(8);
1525
1526 assert!(config.numa_aware);
1527 assert!(config.work_stealing);
1528 assert_eq!(config.num_threads, 8);
1529 }
1530
1531 #[test]
1532 fn test_numa_topology_detection() {
1533 let topology = NumaTopology::detect();
1534
1535 assert!(topology.num_nodes > 0);
1536 assert!(!topology.cores_per_node.is_empty());
1537 assert_eq!(topology.cores_per_node.len(), topology.num_nodes);
1538 assert_eq!(topology.memory_per_node.len(), topology.num_nodes);
1539 }
1540
1541 #[test]
1542 fn test_work_item_creation() {
1543 let item = WorkItem {
1544 start: 0,
1545 end: 100,
1546 work_type: WorkType::DistanceMatrix,
1547 priority: 1,
1548 numa_hint: Some(0),
1549 };
1550
1551 assert_eq!(item.start, 0);
1552 assert_eq!(item.end, 100);
1553 assert_eq!(item.work_type, WorkType::DistanceMatrix);
1554 assert_eq!(item.priority, 1);
1555 assert_eq!(item.numa_hint, Some(0));
1556 }
1557
1558 #[test]
1559 fn test_work_stealing_pool_creation() {
1560 let config = WorkStealingConfig::new().with_threads(1); let pool = WorkStealingPool::new(config);
1562
1563 assert!(pool.is_ok());
1564 let pool = pool.expect("Operation failed");
1565 assert_eq!(pool.workers.len(), 1);
1566 }
1567
1568 #[test]
1569 fn test_advanced_parallel_distance_matrix() {
1570 let _points = array![[0.0, 0.0], [1.0, 0.0]];
1572 let config = WorkStealingConfig::new().with_threads(1);
1573
1574 let processor = AdvancedParallelDistanceMatrix::new(config);
1575 assert!(processor.is_ok());
1576
1577 let processor = processor.expect("Operation failed");
1579 let stats = processor.statistics();
1580 assert_eq!(stats.num_threads, 1);
1581 }
1582
1583 #[test]
1584 fn test_advanced_parallel_kmeans() {
1585 let points = array![[0.0, 0.0], [1.0, 1.0]];
1587 let config = WorkStealingConfig::new().with_threads(1); let kmeans = AdvancedParallelKMeans::new(1, config); assert!(kmeans.is_ok());
1591
1592 let kmeans = kmeans.expect("Operation failed");
1593 let result = kmeans.fit_parallel(&points.view());
1594 assert!(result.is_ok());
1595
1596 let (centroids, assignments) = result.expect("Operation failed");
1597 assert_eq!(centroids.dim(), (1, 2));
1598 assert_eq!(assignments.len(), 2);
1599 }
1600
1601 #[test]
1602 fn test_global_functions() {
1603 let _topology = get_numa_topology();
1605 report_advanced_parallel_capabilities();
1606
1607 let config = WorkStealingConfig::new().with_threads(1);
1608 let init_result = initialize_global_pool(config);
1609 assert!(init_result.is_ok());
1610 }
1611
1612 #[test]
1613 fn test_work_context_structures() {
1614 let (sender, _receiver) = channel::<(usize, usize, f64)>();
1616
1617 let distance_context = DistanceMatrixContext {
1618 points: Array2::zeros((4, 2)),
1619 result_sender: sender,
1620 };
1621
1622 let work_context = WorkContext {
1623 distance_context: Some(distance_context),
1624 kmeans_context: None,
1625 kdtree_context: None,
1626 nn_context: None,
1627 custom_context: None,
1628 };
1629
1630 assert!(work_context.distance_context.is_some());
1632 }
1633
1634 #[test]
1635 fn test_linear_to_matrix_indices() {
1636 let n = 4;
1637 let expected_pairs = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)];
1638
1639 for (_linearidx, expected) in expected_pairs.iter().enumerate() {
1640 let result = WorkStealingPool::linear_to_matrix_indices(_linearidx, n);
1641 assert_eq!(result, *expected, "Failed for linear index {_linearidx}");
1642 }
1643 }
1644
1645 #[test]
1646 fn test_kdtree_chunk_result() {
1647 let chunk_result = KDTreeChunkResult {
1648 node_index: 0,
1649 is_leaf: true,
1650 splitting_dimension: 0,
1651 split_value: 1.0,
1652 left_indices: Vec::new(),
1653 right_indices: Vec::new(),
1654 };
1655
1656 assert!(chunk_result.is_leaf);
1657 assert_eq!(chunk_result.node_index, 0);
1658 assert_eq!(chunk_result.splitting_dimension, 0);
1659 }
1660
1661 #[test]
1662 fn test_enhanced_distance_matrix_computation() {
1663 let _points = array![[0.0, 0.0], [1.0, 0.0]];
1665 let config = WorkStealingConfig::new().with_threads(1);
1666
1667 let processor = AdvancedParallelDistanceMatrix::new(config);
1668 assert!(processor.is_ok());
1669
1670 let processor = processor.expect("Operation failed");
1672 let stats = processor.statistics();
1673 assert_eq!(stats.num_threads, 1);
1674
1675 assert!(stats.numa_nodes > 0);
1679 assert_eq!(stats.numa_nodes, NumaTopology::detect().num_nodes);
1680 }
1681
1682 #[test]
1683 fn test_enhanced_kmeans_with_context() {
1684 let points = array![[0.0, 0.0], [1.0, 1.0]];
1686 let config = WorkStealingConfig::new().with_threads(1); let kmeans = AdvancedParallelKMeans::new(1, config); assert!(kmeans.is_ok());
1690
1691 let kmeans = kmeans.expect("Operation failed");
1692 let result = kmeans.fit_parallel(&points.view());
1693 assert!(result.is_ok());
1694
1695 let (centroids, assignments) = result.expect("Operation failed");
1696 assert_eq!(centroids.dim(), (1, 2));
1697 assert_eq!(assignments.len(), 2);
1698 }
1699
1700 #[test]
1701 fn test_numa_topology_detailed() {
1702 let topology = NumaTopology::detect();
1703
1704 assert!(topology.num_nodes > 0);
1705 assert_eq!(topology.cores_per_node.len(), topology.num_nodes);
1706 assert_eq!(topology.memory_per_node.len(), topology.num_nodes);
1707 assert_eq!(topology.distance_matrix.len(), topology.num_nodes);
1708
1709 for node in 0..topology.num_nodes {
1711 let threads = topology.optimal_threads_per_node(node);
1712 assert!(threads > 0);
1713 }
1714
1715 for node in 0..topology.num_nodes {
1717 let _capacity = topology.memory_capacity(node);
1718 }
1720 }
1721
1722 #[test]
1723 fn test_work_stealing_configuration_advanced() {
1724 let config = WorkStealingConfig::new()
1725 .with_numa_aware(true)
1726 .with_work_stealing(true)
1727 .with_adaptive_scheduling(true)
1728 .with_threads(4)
1729 .with_chunk_sizes(512, 32)
1730 .with_thread_affinity(ThreadAffinityStrategy::NumaAware)
1731 .with_memory_strategy(MemoryStrategy::NumaInterleaved);
1732
1733 assert!(config.numa_aware);
1734 assert!(config.work_stealing);
1735 assert!(config.adaptive_scheduling);
1736 assert_eq!(config.num_threads, 4);
1737 assert_eq!(config.initial_chunk_size, 512);
1738 assert_eq!(config.min_chunk_size, 32);
1739 assert_eq!(config.thread_affinity, ThreadAffinityStrategy::NumaAware);
1740 assert_eq!(config.memory_strategy, MemoryStrategy::NumaInterleaved);
1741 }
1742}