1use crate::error::SpatialResult;
54use crate::memory_pool::DistancePool;
55use scirs2_core::ndarray::{Array1, Array2, ArrayView2};
56use std::path::Path;
57use std::process::Command;
58use std::sync::Arc;
59
60type GpuDeviceInfoResult = Result<(Vec<String>, Vec<(usize, usize)>), Box<dyn std::error::Error>>;
62
63#[derive(Debug, Clone)]
65pub struct GpuCapabilities {
66 pub gpu_available: bool,
68 pub device_count: usize,
70 pub total_memory: usize,
72 pub available_memory: usize,
74 pub compute_capability: Option<(u32, u32)>,
76 pub max_threads_per_block: usize,
78 pub max_blocks_per_grid: usize,
80 pub device_names: Vec<String>,
82 pub supported_backends: Vec<GpuBackend>,
84}
85
86impl Default for GpuCapabilities {
87 fn default() -> Self {
88 Self {
89 gpu_available: false,
90 device_count: 0,
91 total_memory: 0,
92 available_memory: 0,
93 compute_capability: None,
94 max_threads_per_block: 1024,
95 max_blocks_per_grid: 65535,
96 device_names: Vec::new(),
97 supported_backends: Vec::new(),
98 }
99 }
100}
101
102#[derive(Debug, Clone, PartialEq)]
104pub enum GpuBackend {
105 Cuda,
107 Rocm,
109 LevelZero,
111 Vulkan,
113 CpuFallback,
115}
116
117pub struct GpuDevice {
119 capabilities: GpuCapabilities,
120 preferred_backend: GpuBackend,
121 #[allow(dead_code)]
122 memory_pool: Arc<DistancePool>,
123}
124
125impl GpuDevice {
126 pub fn new() -> SpatialResult<Self> {
128 let capabilities = Self::detect_capabilities()?;
129 let preferred_backend = Self::select_optimal_backend(&capabilities);
130 let memory_pool = Arc::new(DistancePool::new(1000));
131
132 Ok(Self {
133 capabilities,
134 preferred_backend,
135 memory_pool,
136 })
137 }
138
139 fn detect_capabilities() -> SpatialResult<GpuCapabilities> {
141 let mut caps = GpuCapabilities::default();
142
143 #[cfg(feature = "cuda")]
145 {
146 if Self::check_cuda_available() {
147 caps.gpu_available = true;
148 caps.device_count = Self::get_cuda_device_count();
149 caps.supported_backends.push(GpuBackend::Cuda);
150
151 if let Ok((names, memory_info)) = Self::get_cuda_device_info() {
153 caps.device_names = names;
154 if let Some((total, available)) = memory_info.first() {
155 caps.total_memory = *total;
156 caps.available_memory = *available;
157 }
158 }
159
160 caps.max_threads_per_block = 1024;
162 caps.max_blocks_per_grid = 2147483647; caps.compute_capability = Self::get_cuda_compute_capability();
164 }
165 }
166
167 #[cfg(feature = "rocm")]
169 {
170 if Self::check_rocm_available() {
171 caps.gpu_available = true;
172 let rocm_count = Self::get_rocm_device_count();
173 if rocm_count > caps.device_count {
174 caps.device_count = rocm_count;
175 }
176 caps.supported_backends.push(GpuBackend::Rocm);
177
178 if let Ok((names, memory_info)) = Self::get_rocm_device_info() {
180 if caps.device_names.is_empty() {
181 caps.device_names = names;
182 } else {
183 caps.device_names.extend(names);
184 }
185 if let Some((total, available)) = memory_info.first() {
186 if caps.total_memory == 0 {
187 caps.total_memory = *total;
188 caps.available_memory = *available;
189 }
190 }
191 }
192
193 caps.max_threads_per_block = 1024;
195 caps.max_blocks_per_grid = 2147483647;
196 }
197 }
198
199 #[cfg(feature = "vulkan")]
201 {
202 if Self::check_vulkan_available() {
203 caps.gpu_available = true;
204 caps.supported_backends.push(GpuBackend::Vulkan);
205
206 if let Ok((names, memory_info)) = Self::get_vulkan_device_info() {
208 if caps.device_names.is_empty() {
209 caps.device_names = names;
210 } else {
211 caps.device_names.extend(names);
212 }
213 if let Some((total, available)) = memory_info.first() {
214 if caps.total_memory == 0 {
215 caps.total_memory = *total;
216 caps.available_memory = *available;
217 }
218 }
219 }
220 }
221 }
222
223 caps.supported_backends.push(GpuBackend::CpuFallback);
225
226 Ok(caps)
227 }
228
229 fn select_optimal_backend(caps: &GpuCapabilities) -> GpuBackend {
231 if caps.supported_backends.contains(&GpuBackend::Cuda) {
233 GpuBackend::Cuda
234 } else if caps.supported_backends.contains(&GpuBackend::Rocm) {
235 GpuBackend::Rocm
236 } else if caps.supported_backends.contains(&GpuBackend::Vulkan) {
237 GpuBackend::Vulkan
238 } else {
239 GpuBackend::CpuFallback
240 }
241 }
242
243 pub fn is_gpu_available(&self) -> bool {
245 self.capabilities.gpu_available
246 }
247
248 pub fn capabilities(&self) -> &GpuCapabilities {
250 &self.capabilities
251 }
252
253 pub fn optimal_block_size(&self, _problemsize: usize) -> usize {
255 match self.preferred_backend {
256 GpuBackend::Cuda => {
257 let warp_size = 32;
259 let optimal = (_problemsize / warp_size).max(1) * warp_size;
260 optimal.min(self.capabilities.max_threads_per_block)
261 }
262 GpuBackend::Rocm => {
263 let wavefront_size = 64;
265 let optimal = (_problemsize / wavefront_size).max(1) * wavefront_size;
266 optimal.min(self.capabilities.max_threads_per_block)
267 }
268 _ => {
269 256.min(self.capabilities.max_threads_per_block)
271 }
272 }
273 }
274
275 #[cfg(feature = "cuda")]
277 fn check_cuda_available() -> bool {
278 if let Ok(output) = Command::new("nvidia-smi")
280 .arg("--query-gpu=count")
281 .arg("--format=csv,noheader,nounits")
282 .output()
283 {
284 if output.status.success() {
285 if let Ok(count_str) = String::from_utf8(output.stdout) {
286 if let Ok(count) = count_str.trim().parse::<u32>() {
287 return count > 0;
288 }
289 }
290 }
291 }
292
293 #[cfg(target_os = "linux")]
295 {
296 Path::exists(Path::new("/usr/local/cuda/lib64/libcuda.so"))
297 || Path::exists(Path::new("/usr/lib/x86_64-linux-gnu/libcuda.so"))
298 || Path::exists(Path::new("/usr/lib64/libcuda.so"))
299 }
300
301 #[cfg(target_os = "windows")]
302 {
303 Path::exists(Path::new("C:\\Windows\\System32\\nvcuda.dll"))
304 }
305
306 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
307 false
308 }
309
310 #[cfg(feature = "cuda")]
311 fn get_cuda_device_count() -> usize {
312 if let Ok(output) = Command::new("nvidia-smi")
313 .arg("--query-gpu=count")
314 .arg("--format=csv,noheader,nounits")
315 .output()
316 {
317 if output.status.success() {
318 if let Ok(count_str) = String::from_utf8(output.stdout) {
319 if let Ok(count) = count_str.trim().parse::<usize>() {
320 return count;
321 }
322 }
323 }
324 }
325
326 if let Ok(output) = Command::new("nvidia-smi").arg("-L").output() {
328 if output.status.success() {
329 if let Ok(list_str) = String::from_utf8(output.stdout) {
330 return list_str
331 .lines()
332 .filter(|line| line.starts_with("GPU "))
333 .count();
334 }
335 }
336 }
337
338 0
339 }
340
341 #[cfg(feature = "rocm")]
342 fn check_rocm_available() -> bool {
343 if let Ok(output) = Command::new("rocm-smi").arg("--showid").output() {
345 if output.status.success() {
346 return true;
347 }
348 }
349
350 #[cfg(target_os = "linux")]
352 {
353 Path::exists(Path::new("/opt/rocm/lib/libhip.so"))
354 || Path::exists(Path::new("/usr/lib/libhip.so"))
355 || Path::exists(Path::new("/usr/lib/x86_64-linux-gnu/libhip.so"))
356 }
357
358 #[cfg(not(target_os = "linux"))]
359 false
360 }
361
362 #[cfg(feature = "rocm")]
363 fn get_rocm_device_count() -> usize {
364 if let Ok(output) = Command::new("rocm-smi").arg("--showid").output() {
365 if output.status.success() {
366 if let Ok(list_str) = String::from_utf8(output.stdout) {
367 return list_str
369 .lines()
370 .filter(|line| line.contains("GPU") || line.contains("card"))
371 .count();
372 }
373 }
374 }
375
376 #[cfg(target_os = "linux")]
378 {
379 use std::fs;
380 if let Ok(entries) = fs::read_dir("/sys/class/drm") {
381 let count = entries
382 .filter_map(Result::ok)
383 .filter(|entry| {
384 if let Ok(name) = entry.file_name().into_string() {
385 name.starts_with("card") && !name.contains("-")
386 } else {
387 false
388 }
389 })
390 .count();
391 if count > 0 {
392 return count;
393 }
394 }
395 }
396
397 0
398 }
399
400 #[cfg(feature = "vulkan")]
401 fn check_vulkan_available() -> bool {
402 if let Ok(output) = Command::new("vulkaninfo").arg("--summary").output() {
404 if output.status.success() {
405 if let Ok(info_str) = String::from_utf8(output.stdout) {
406 return info_str.contains("VK_QUEUE_COMPUTE_BIT")
408 || info_str.contains("deviceType");
409 }
410 }
411 }
412
413 #[cfg(target_os = "linux")]
415 {
416 Path::exists(Path::new("/usr/lib/libvulkan.so"))
417 || Path::exists(Path::new("/usr/lib/x86_64-linux-gnu/libvulkan.so"))
418 || Path::exists(Path::new("/usr/local/lib/libvulkan.so"))
419 }
420
421 #[cfg(target_os = "windows")]
422 {
423 Path::exists(Path::new("C:\\Windows\\System32\\vulkan-1.dll"))
424 }
425
426 #[cfg(target_os = "macos")]
427 {
428 Path::exists(Path::new("/usr/local/lib/libvulkan.dylib"))
429 || Path::exists(Path::new(
430 "/System/Library/Frameworks/Metal.framework/Metal",
431 ))
432 }
433
434 #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos")))]
435 false
436 }
437
438 #[cfg(feature = "cuda")]
440 fn get_cuda_device_info() -> GpuDeviceInfoResult {
441 let mut device_names = Vec::new();
442 let mut memory_info = Vec::new();
443
444 if let Ok(output) = Command::new("nvidia-smi")
446 .arg("--query-gpu=name")
447 .arg("--format=csv,noheader,nounits")
448 .output()
449 {
450 if output.status.success() {
451 if let Ok(names_str) = String::from_utf8(output.stdout) {
452 device_names = names_str.lines().map(|s| s.trim().to_string()).collect();
453 }
454 }
455 }
456
457 if let Ok(output) = Command::new("nvidia-smi")
459 .arg("--query-gpu=memory.total,memory.free")
460 .arg("--format=csv,noheader,nounits")
461 .output()
462 {
463 if output.status.success() {
464 if let Ok(memory_str) = String::from_utf8(output.stdout) {
465 for line in memory_str.lines() {
466 let parts: Vec<&str> = line.split(',').collect();
467 if parts.len() >= 2 {
468 if let (Ok(total), Ok(free)) = (
469 parts[0].trim().parse::<usize>(),
470 parts[1].trim().parse::<usize>(),
471 ) {
472 memory_info.push((total * 1024 * 1024, free * 1024 * 1024));
473 }
475 }
476 }
477 }
478 }
479 }
480
481 Ok((device_names, memory_info))
482 }
483
484 #[cfg(feature = "cuda")]
486 fn get_cuda_compute_capability() -> Option<(u32, u32)> {
487 if let Ok(output) = Command::new("nvidia-smi")
488 .arg("--query-gpu=compute_cap")
489 .arg("--format=csv,noheader,nounits")
490 .output()
491 {
492 if output.status.success() {
493 if let Ok(cap_str) = String::from_utf8(output.stdout) {
494 if let Some(line) = cap_str.lines().next() {
495 let parts: Vec<&str> = line.trim().split('.').collect();
496 if parts.len() >= 2 {
497 if let (Ok(major), Ok(minor)) =
498 (parts[0].parse::<u32>(), parts[1].parse::<u32>())
499 {
500 return Some((major, minor));
501 }
502 }
503 }
504 }
505 }
506 }
507
508 None
509 }
510
511 #[cfg(feature = "rocm")]
513 fn get_rocm_device_info() -> GpuDeviceInfoResult {
514 let mut device_names = Vec::new();
515 let mut memory_info = Vec::new();
516
517 if let Ok(output) = Command::new("rocm-smi").arg("--showproductname").output() {
519 if output.status.success() {
520 if let Ok(info_str) = String::from_utf8(output.stdout) {
521 for line in info_str.lines() {
522 if line.contains("Card series:") {
523 if let Some(name) = line.split(':').nth(1) {
524 device_names.push(name.trim().to_string());
525 }
526 }
527 }
528 }
529 }
530 }
531
532 if let Ok(output) = Command::new("rocm-smi")
534 .arg("--showmeminfo")
535 .arg("vram")
536 .output()
537 {
538 if output.status.success() {
539 if let Ok(memory_str) = String::from_utf8(output.stdout) {
540 for line in memory_str.lines() {
541 if line.contains("Total memory") || line.contains("Used memory") {
542 if let Some(mem_part) = line
543 .split_whitespace()
544 .find(|s| s.ends_with("MB") || s.ends_with("GB"))
545 {
546 if let Ok(mem_val) = mem_part
547 .trim_end_matches("MB")
548 .trim_end_matches("GB")
549 .parse::<usize>()
550 {
551 let bytes = if mem_part.ends_with("GB") {
552 mem_val * 1024 * 1024 * 1024
553 } else {
554 mem_val * 1024 * 1024
555 };
556 memory_info.push((bytes, bytes / 2));
557 }
558 }
559 }
560 }
561 }
562 }
563 }
564
565 if device_names.is_empty() && memory_info.is_empty() {
566 device_names.push("AMD GPU (ROCm)".to_string());
567 memory_info.push((8 * 1024 * 1024 * 1024, 6 * 1024 * 1024 * 1024));
568 }
569
570 Ok((device_names, memory_info))
571 }
572
573 #[cfg(feature = "vulkan")]
575 fn get_vulkan_device_info() -> GpuDeviceInfoResult {
576 let mut device_names = Vec::new();
577 let mut memory_info = Vec::new();
578
579 if let Ok(output) = Command::new("vulkaninfo").arg("--summary").output() {
580 if output.status.success() {
581 if let Ok(info_str) = String::from_utf8(output.stdout) {
582 for line in info_str.lines() {
583 if line.contains("deviceName") {
584 if let Some(name_part) = line.split('=').nth(1) {
585 device_names.push(name_part.trim().to_string());
586 }
587 } else if line.contains("heapSize") {
588 if let Some(mem_part) = line.split('=').nth(1) {
589 if let Ok(mem_val) = mem_part.trim().parse::<usize>() {
590 memory_info.push((mem_val, mem_val * 3 / 4));
591 }
592 }
593 }
594 }
595 }
596 }
597 }
598
599 if device_names.is_empty() {
600 device_names.push("Vulkan Device".to_string());
601 memory_info.push((4 * 1024 * 1024 * 1024, 3 * 1024 * 1024 * 1024));
602 }
603
604 Ok((device_names, memory_info))
605 }
606
607 #[cfg(not(feature = "cuda"))]
608 #[allow(dead_code)]
609 fn get_cuda_device_info() -> GpuDeviceInfoResult {
610 Ok((Vec::new(), Vec::new()))
611 }
612
613 #[cfg(not(feature = "cuda"))]
614 #[allow(dead_code)]
615 fn get_cuda_compute_capability() -> Option<(u32, u32)> {
616 None
617 }
618
619 #[cfg(not(feature = "rocm"))]
620 #[allow(dead_code)]
621 fn get_rocm_device_info() -> GpuDeviceInfoResult {
622 Ok((Vec::new(), Vec::new()))
623 }
624
625 #[cfg(not(feature = "vulkan"))]
626 #[allow(dead_code)]
627 fn get_vulkan_device_info() -> GpuDeviceInfoResult {
628 Ok((Vec::new(), Vec::new()))
629 }
630}
631
632impl Default for GpuDevice {
633 fn default() -> Self {
634 Self::new().unwrap_or_else(|_| Self {
635 capabilities: GpuCapabilities::default(),
636 preferred_backend: GpuBackend::CpuFallback,
637 memory_pool: Arc::new(DistancePool::new(1000)),
638 })
639 }
640}
641
642pub struct GpuDistanceMatrix {
644 device: Arc<GpuDevice>,
645 batch_size: usize,
646 use_mixed_precision: bool,
647}
648
649impl GpuDistanceMatrix {
650 pub fn new() -> SpatialResult<Self> {
652 let device = Arc::new(GpuDevice::new()?);
653 Ok(Self {
654 device,
655 batch_size: 1024,
656 use_mixed_precision: true,
657 })
658 }
659
660 pub fn with_batch_size(mut self, batchsize: usize) -> Self {
662 self.batch_size = batchsize;
663 self
664 }
665
666 pub fn with_mixed_precision(mut self, use_mixedprecision: bool) -> Self {
668 self.use_mixed_precision = use_mixedprecision;
669 self
670 }
671
672 pub async fn compute_parallel(
675 &self,
676 points: &ArrayView2<'_, f64>,
677 ) -> SpatialResult<Array2<f64>> {
678 let _n_points = points.nrows();
679
680 if !self.device.is_gpu_available() {
681 return self.compute_cpu_fallback(points).await;
682 }
683
684 match self.device.preferred_backend {
685 GpuBackend::Cuda => self.compute_cuda(points).await,
686 GpuBackend::Rocm => self.compute_rocm(points).await,
687 GpuBackend::Vulkan => self.compute_vulkan(points).await,
688 GpuBackend::CpuFallback => self.compute_cpu_fallback(points).await,
689 GpuBackend::LevelZero => self.compute_cpu_fallback(points).await, }
691 }
692
693 async fn compute_cuda(&self, points: &ArrayView2<'_, f64>) -> SpatialResult<Array2<f64>> {
695 self.compute_cpu_fallback(points).await
703 }
704
705 async fn compute_rocm(&self, points: &ArrayView2<'_, f64>) -> SpatialResult<Array2<f64>> {
707 self.compute_cpu_fallback(points).await
709 }
710
711 async fn compute_vulkan(&self, points: &ArrayView2<'_, f64>) -> SpatialResult<Array2<f64>> {
713 self.compute_cpu_fallback(points).await
715 }
716
717 async fn compute_cpu_fallback(
719 &self,
720 points: &ArrayView2<'_, f64>,
721 ) -> SpatialResult<Array2<f64>> {
722 use crate::simd_distance::parallel_pdist;
724
725 let condensed = parallel_pdist(points, "euclidean")?;
726
727 let n = points.nrows();
729 let mut matrix = Array2::zeros((n, n));
730
731 let mut idx = 0;
732 for i in 0..n {
733 for j in (i + 1)..n {
734 matrix[[i, j]] = condensed[idx];
735 matrix[[j, i]] = condensed[idx];
736 idx += 1;
737 }
738 }
739
740 Ok(matrix)
741 }
742}
743
744pub struct GpuKMeans {
746 device: Arc<GpuDevice>,
747 k: usize,
748 max_iterations: usize,
749 tolerance: f64,
750 batch_size: usize,
751}
752
753impl GpuKMeans {
754 pub fn new(k: usize) -> SpatialResult<Self> {
756 let device = Arc::new(GpuDevice::new()?);
757 Ok(Self {
758 device,
759 k,
760 max_iterations: 100,
761 tolerance: 1e-6,
762 batch_size: 1024,
763 })
764 }
765
766 pub fn with_max_iterations(mut self, maxiterations: usize) -> Self {
768 self.max_iterations = maxiterations;
769 self
770 }
771
772 pub fn with_tolerance(mut self, tolerance: f64) -> Self {
774 self.tolerance = tolerance;
775 self
776 }
777
778 pub fn with_batch_size(mut self, batchsize: usize) -> Self {
780 self.batch_size = batchsize;
781 self
782 }
783
784 pub async fn fit(
787 &self,
788 points: &ArrayView2<'_, f64>,
789 ) -> SpatialResult<(Array2<f64>, Array1<usize>)> {
790 if !self.device.is_gpu_available() {
791 return self.fit_cpu_fallback(points).await;
792 }
793
794 match self.device.preferred_backend {
795 GpuBackend::Cuda => self.fit_cuda(points).await,
796 GpuBackend::Rocm => self.fit_rocm(points).await,
797 GpuBackend::Vulkan => self.fit_vulkan(points).await,
798 GpuBackend::CpuFallback => self.fit_cpu_fallback(points).await,
799 GpuBackend::LevelZero => self.fit_cpu_fallback(points).await, }
801 }
802
803 async fn fit_cuda(
805 &self,
806 points: &ArrayView2<'_, f64>,
807 ) -> SpatialResult<(Array2<f64>, Array1<usize>)> {
808 self.fit_cpu_fallback(points).await
815 }
816
817 async fn fit_rocm(
819 &self,
820 points: &ArrayView2<'_, f64>,
821 ) -> SpatialResult<(Array2<f64>, Array1<usize>)> {
822 self.fit_cpu_fallback(points).await
824 }
825
826 async fn fit_vulkan(
828 &self,
829 points: &ArrayView2<'_, f64>,
830 ) -> SpatialResult<(Array2<f64>, Array1<usize>)> {
831 self.fit_cpu_fallback(points).await
833 }
834
835 async fn fit_cpu_fallback(
837 &self,
838 points: &ArrayView2<'_, f64>,
839 ) -> SpatialResult<(Array2<f64>, Array1<usize>)> {
840 use crate::simd_distance::advanced_simd_clustering::AdvancedSimdKMeans;
842
843 let advanced_kmeans = AdvancedSimdKMeans::new(self.k)
844 .with_mixed_precision(true)
845 .with_block_size(256);
846
847 advanced_kmeans.fit(points)
848 }
849}
850
851pub struct GpuNearestNeighbors {
853 device: Arc<GpuDevice>,
854 #[allow(dead_code)]
855 build_batch_size: usize,
856 #[allow(dead_code)]
857 query_batch_size: usize,
858}
859
860impl GpuNearestNeighbors {
861 pub fn new() -> SpatialResult<Self> {
863 let device = Arc::new(GpuDevice::new()?);
864 Ok(Self {
865 device,
866 build_batch_size: 1024,
867 query_batch_size: 256,
868 })
869 }
870
871 pub async fn knn_search(
874 &self,
875 query_points: &ArrayView2<'_, f64>,
876 data_points: &ArrayView2<'_, f64>,
877 k: usize,
878 ) -> SpatialResult<(Array2<usize>, Array2<f64>)> {
879 if !self.device.is_gpu_available() {
880 return self
881 .knn_search_cpu_fallback(query_points, data_points, k)
882 .await;
883 }
884
885 match self.device.preferred_backend {
886 GpuBackend::Cuda => self.knn_search_cuda(query_points, data_points, k).await,
887 GpuBackend::Rocm => self.knn_search_rocm(query_points, data_points, k).await,
888 GpuBackend::Vulkan => self.knn_search_vulkan(query_points, data_points, k).await,
889 GpuBackend::CpuFallback => {
890 self.knn_search_cpu_fallback(query_points, data_points, k)
891 .await
892 }
893 _ => {
894 self.knn_search_cpu_fallback(query_points, data_points, k)
895 .await
896 }
897 }
898 }
899
900 async fn knn_search_cuda(
902 &self,
903 query_points: &ArrayView2<'_, f64>,
904 data_points: &ArrayView2<'_, f64>,
905 k: usize,
906 ) -> SpatialResult<(Array2<usize>, Array2<f64>)> {
907 self.knn_search_cpu_fallback(query_points, data_points, k)
913 .await
914 }
915
916 async fn knn_search_rocm(
918 &self,
919 query_points: &ArrayView2<'_, f64>,
920 data_points: &ArrayView2<'_, f64>,
921 k: usize,
922 ) -> SpatialResult<(Array2<usize>, Array2<f64>)> {
923 self.knn_search_cpu_fallback(query_points, data_points, k)
924 .await
925 }
926
927 async fn knn_search_vulkan(
929 &self,
930 query_points: &ArrayView2<'_, f64>,
931 data_points: &ArrayView2<'_, f64>,
932 k: usize,
933 ) -> SpatialResult<(Array2<usize>, Array2<f64>)> {
934 self.knn_search_cpu_fallback(query_points, data_points, k)
935 .await
936 }
937
938 async fn knn_search_cpu_fallback(
940 &self,
941 query_points: &ArrayView2<'_, f64>,
942 data_points: &ArrayView2<'_, f64>,
943 k: usize,
944 ) -> SpatialResult<(Array2<usize>, Array2<f64>)> {
945 use crate::simd_distance::advanced_simd_clustering::AdvancedSimdNearestNeighbors;
947
948 let advanced_nn = AdvancedSimdNearestNeighbors::new();
949 advanced_nn.simd_knn_advanced_fast(query_points, data_points, k)
950 }
951}
952
953impl Default for GpuNearestNeighbors {
954 fn default() -> Self {
955 Self::new().unwrap_or_else(|_| Self {
956 device: Arc::new(GpuDevice::default()),
957 build_batch_size: 1024,
958 query_batch_size: 256,
959 })
960 }
961}
962
963pub struct HybridProcessor {
965 gpu_device: Arc<GpuDevice>,
966 cpu_threshold: usize,
967 gpu_threshold: usize,
968}
969
970impl HybridProcessor {
971 pub fn new() -> SpatialResult<Self> {
973 let gpu_device = Arc::new(GpuDevice::new()?);
974 Ok(Self {
975 gpu_device,
976 cpu_threshold: 1000, gpu_threshold: 100000, })
979 }
980
981 pub fn choose_strategy(&self, _datasetsize: usize) -> ProcessingStrategy {
983 if !self.gpu_device.is_gpu_available() {
984 return ProcessingStrategy::CpuOnly;
985 }
986
987 if _datasetsize < self.cpu_threshold {
988 ProcessingStrategy::CpuOnly
989 } else if _datasetsize < self.gpu_threshold {
990 ProcessingStrategy::Hybrid
991 } else {
992 ProcessingStrategy::GpuOnly
993 }
994 }
995
996 pub fn optimal_batch_sizes(&self, _totalsize: usize) -> (usize, usize) {
998 let gpu_capability = self.gpu_device.capabilities().total_memory / (8 * 1024); let cpu_batch = (_totalsize / 4).max(1000); let gpu_batch = (_totalsize * 3 / 4).min(gpu_capability); (cpu_batch, gpu_batch)
1003 }
1004}
1005
1006impl Default for HybridProcessor {
1007 fn default() -> Self {
1008 Self::new().unwrap_or_else(|_| Self {
1009 gpu_device: Arc::new(GpuDevice::default()),
1010 cpu_threshold: 1000,
1011 gpu_threshold: 100000,
1012 })
1013 }
1014}
1015
1016#[derive(Debug, Clone, PartialEq)]
1018pub enum ProcessingStrategy {
1019 CpuOnly,
1021 GpuOnly,
1023 Hybrid,
1025}
1026
1027static GLOBAL_GPU_DEVICE: std::sync::OnceLock<GpuDevice> = std::sync::OnceLock::new();
1029
1030#[allow(dead_code)]
1032pub fn global_gpu_device() -> &'static GpuDevice {
1033 GLOBAL_GPU_DEVICE.get_or_init(GpuDevice::default)
1034}
1035
1036#[allow(dead_code)]
1038pub fn is_gpu_acceleration_available() -> bool {
1039 global_gpu_device().is_gpu_available()
1040}
1041
1042#[allow(dead_code)]
1044pub fn get_gpu_capabilities() -> &'static GpuCapabilities {
1045 global_gpu_device().capabilities()
1046}
1047
1048#[allow(dead_code)]
1050pub fn report_gpu_status() {
1051 let device = global_gpu_device();
1052 let caps = device.capabilities();
1053
1054 println!("GPU Acceleration Status:");
1055 println!(" Available: {}", caps.gpu_available);
1056 println!(" Device Count: {}", caps.device_count);
1057
1058 if caps.gpu_available {
1059 println!(
1060 " Total Memory: {:.1} GB",
1061 caps.total_memory as f64 / (1024.0 * 1024.0 * 1024.0)
1062 );
1063 println!(
1064 " Available Memory: {:.1} GB",
1065 caps.available_memory as f64 / (1024.0 * 1024.0 * 1024.0)
1066 );
1067 println!(" Max Threads/Block: {}", caps.max_threads_per_block);
1068 println!(" Supported Backends: {:?}", caps.supported_backends);
1069
1070 for (i, name) in caps.device_names.iter().enumerate() {
1071 println!(" Device {i}: {name}");
1072 }
1073 } else {
1074 println!(" Reason: No compatible GPU devices found");
1075 println!(" Fallback: Using optimized CPU SIMD operations");
1076 }
1077}
1078
1079#[cfg(test)]
1080mod tests {
1081 use super::*;
1082 use scirs2_core::ndarray::array;
1083
1084 #[test]
1085 fn test_gpu_device_creation() {
1086 let device = GpuDevice::new();
1087 assert!(device.is_ok());
1088
1089 let device = device.expect("Operation failed");
1090 assert!(!device.capabilities().supported_backends.is_empty());
1092 }
1093
1094 #[test]
1095 fn test_processing_strategy_selection() {
1096 let processor = HybridProcessor::new().expect("Operation failed");
1097
1098 let strategy = processor.choose_strategy(500);
1100 assert_eq!(strategy, ProcessingStrategy::CpuOnly);
1101
1102 let strategy = processor.choose_strategy(200000);
1104 assert!(matches!(
1106 strategy,
1107 ProcessingStrategy::GpuOnly | ProcessingStrategy::CpuOnly
1108 ));
1109 }
1110
1111 #[cfg(feature = "async")]
1112 #[tokio::test]
1113 async fn test_gpu_distance_matrix() {
1114 let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]];
1115
1116 let gpu_matrix = GpuDistanceMatrix::new().expect("Operation failed");
1117 let points_view = points.view();
1118 let result = gpu_matrix.compute_parallel(&points_view).await;
1119
1120 assert!(result.is_ok());
1121 let matrix = result.expect("Operation failed");
1122 assert_eq!(matrix.dim(), (4, 4));
1123 }
1124
1125 #[cfg(feature = "async")]
1126 #[tokio::test]
1127 async fn test_gpu_kmeans() {
1128 let points = array![
1129 [0.0, 0.0],
1130 [0.1, 0.1],
1131 [0.0, 0.1], [5.0, 5.0],
1133 [5.1, 5.1],
1134 [5.0, 5.1], ];
1136
1137 let gpu_kmeans = GpuKMeans::new(2).expect("Operation failed");
1138 let points_view = points.view();
1139 let result = gpu_kmeans.fit(&points_view).await;
1140
1141 assert!(result.is_ok());
1142 let (centroids, _assignments) = result.expect("Operation failed");
1143 assert_eq!(centroids.nrows(), 2);
1144 }
1145
1146 #[cfg(feature = "async")]
1147 #[tokio::test]
1148 async fn test_gpu_nearest_neighbors() {
1149 let data_points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]];
1150 let query_points = array![[0.1, 0.1], [0.9, 0.9]];
1151
1152 let gpu_nn = GpuNearestNeighbors::new().expect("Operation failed");
1153 let query_view = query_points.view();
1154 let data_view = data_points.view();
1155 let result = gpu_nn.knn_search(&query_view, &data_view, 2).await;
1156
1157 assert!(result.is_ok());
1158 let (indices, _distances) = result.expect("Operation failed");
1159 assert_eq!(indices[[0, 0]], 0); }
1162
1163 #[test]
1164 fn test_global_gpu_functions() {
1165 let device = global_gpu_device();
1167 assert!(!device.capabilities.device_names.is_empty() || !device.capabilities.gpu_available);
1168
1169 report_gpu_status();
1171 let _caps = get_gpu_capabilities();
1172 let _available = is_gpu_acceleration_available();
1173 }
1174}