1use std::collections::HashMap;
7use std::sync::{Arc, RwLock};
8use std::time::Instant;
9
10#[derive(Debug, Clone)]
12pub struct GpuDevice {
13 pub id: u32,
14 pub name: String,
15 pub memory_total: u64,
16 pub memory_available: u64,
17 pub compute_capability: (u32, u32),
18 pub cores: u32,
19 pub clock_rate: u32,
20 pub memory_bandwidth: u64,
21 pub is_integrated: bool,
22}
23
24#[derive(Debug, Clone)]
26pub struct GpuMemoryAllocation {
27 pub ptr: u64,
28 pub size: u64,
29 pub device_id: u32,
30 pub allocated_at: Instant,
31 pub name: String,
32}
33
34#[derive(Debug, Clone)]
36pub struct GpuKernelExecution {
37 pub kernel_name: String,
38 pub device_id: u32,
39 pub grid_size: (u32, u32, u32),
40 pub block_size: (u32, u32, u32),
41 pub shared_memory: u32,
42 pub execution_time: f64,
43 pub parameters: HashMap<String, String>,
44}
45
46#[derive(Debug)]
48pub struct GpuUtils {
49 devices: Vec<GpuDevice>,
50 allocations: Arc<RwLock<HashMap<u64, GpuMemoryAllocation>>>,
51 kernel_executions: Arc<RwLock<Vec<GpuKernelExecution>>>,
52 performance_counters: Arc<RwLock<HashMap<String, f64>>>,
53}
54
55impl GpuUtils {
56 pub fn new() -> Self {
58 Self {
59 devices: Vec::new(),
60 allocations: Arc::new(RwLock::new(HashMap::new())),
61 kernel_executions: Arc::new(RwLock::new(Vec::new())),
62 performance_counters: Arc::new(RwLock::new(HashMap::new())),
63 }
64 }
65
66 pub fn init_devices(&mut self) -> Result<(), GpuError> {
68 let mock_devices = vec![
70 GpuDevice {
71 id: 0,
72 name: "NVIDIA GeForce RTX 3080".to_string(),
73 memory_total: 10_737_418_240, memory_available: 9_663_676_416, compute_capability: (8, 6),
76 cores: 8704,
77 clock_rate: 1710,
78 memory_bandwidth: 760_000_000_000, is_integrated: false,
80 },
81 GpuDevice {
82 id: 1,
83 name: "Intel UHD Graphics 770".to_string(),
84 memory_total: 2_147_483_648, memory_available: 1_610_612_736, compute_capability: (0, 0),
87 cores: 256,
88 clock_rate: 1550,
89 memory_bandwidth: 68_000_000_000, is_integrated: true,
91 },
92 ];
93
94 self.devices = mock_devices;
95 Ok(())
96 }
97
98 pub fn get_devices(&self) -> &[GpuDevice] {
100 &self.devices
101 }
102
103 pub fn get_device(&self, id: u32) -> Option<&GpuDevice> {
105 self.devices.iter().find(|d| d.id == id)
106 }
107
108 pub fn get_best_device(&self) -> Option<&GpuDevice> {
110 self.devices
111 .iter()
112 .filter(|d| !d.is_integrated)
113 .max_by_key(|d| d.cores * d.clock_rate)
114 .or_else(|| self.devices.first())
115 }
116
117 pub fn allocate_memory(&self, size: u64, device_id: u32, name: &str) -> Result<u64, GpuError> {
119 let device = self.get_device(device_id).ok_or(GpuError::DeviceNotFound)?;
120
121 if size > device.memory_available {
122 return Err(GpuError::OutOfMemory);
123 }
124
125 let ptr = (std::ptr::null::<u8>() as u64) + size; let allocation = GpuMemoryAllocation {
128 ptr,
129 size,
130 device_id,
131 allocated_at: Instant::now(),
132 name: name.to_string(),
133 };
134
135 self.allocations
136 .write()
137 .expect("operation should succeed")
138 .insert(ptr, allocation);
139 Ok(ptr)
140 }
141
142 pub fn free_memory(&self, ptr: u64) -> Result<(), GpuError> {
144 let mut allocations = self.allocations.write().expect("operation should succeed");
145 allocations.remove(&ptr).ok_or(GpuError::InvalidPointer)?;
146 Ok(())
147 }
148
149 pub fn get_memory_stats(&self) -> HashMap<u32, MemoryStats> {
151 let allocations = self.allocations.read().expect("operation should succeed");
152 let mut stats = HashMap::new();
153
154 for device in &self.devices {
155 let device_allocations: Vec<_> = allocations
156 .values()
157 .filter(|a| a.device_id == device.id)
158 .collect();
159
160 let total_allocated = device_allocations.iter().map(|a| a.size).sum();
161 let num_allocations = device_allocations.len();
162
163 stats.insert(
164 device.id,
165 MemoryStats {
166 total_memory: device.memory_total,
167 available_memory: device.memory_available,
168 allocated_memory: total_allocated,
169 free_memory: device.memory_available - total_allocated,
170 num_allocations,
171 largest_allocation: device_allocations
172 .iter()
173 .map(|a| a.size)
174 .max()
175 .unwrap_or(0),
176 fragmentation_ratio: if num_allocations > 0 {
177 (num_allocations as f64) / (total_allocated as f64 / 1024.0)
178 } else {
179 0.0
180 },
181 },
182 );
183 }
184
185 stats
186 }
187
188 pub fn execute_kernel(&self, kernel: &GpuKernelInfo) -> Result<GpuKernelExecution, GpuError> {
190 let _device = self
191 .get_device(kernel.device_id)
192 .ok_or(GpuError::DeviceNotFound)?;
193
194 let start_time = Instant::now();
195
196 std::thread::sleep(std::time::Duration::from_millis(1));
198
199 let execution_time = start_time.elapsed().as_secs_f64() * 1000.0; let execution = GpuKernelExecution {
202 kernel_name: kernel.name.clone(),
203 device_id: kernel.device_id,
204 grid_size: kernel.grid_size,
205 block_size: kernel.block_size,
206 shared_memory: kernel.shared_memory,
207 execution_time,
208 parameters: kernel.parameters.clone(),
209 };
210
211 self.kernel_executions
212 .write()
213 .expect("operation should succeed")
214 .push(execution.clone());
215 Ok(execution)
216 }
217
218 pub fn get_kernel_history(&self) -> Vec<GpuKernelExecution> {
220 self.kernel_executions
221 .read()
222 .expect("operation should succeed")
223 .clone()
224 }
225
226 pub fn get_performance_counters(&self) -> HashMap<String, f64> {
228 self.performance_counters
229 .read()
230 .expect("operation should succeed")
231 .clone()
232 }
233
234 pub fn update_counter(&self, name: &str, value: f64) {
236 self.performance_counters
237 .write()
238 .expect("operation should succeed")
239 .insert(name.to_string(), value);
240 }
241
242 pub fn estimate_throughput(&self, device_id: u32, array_size: usize, operation: &str) -> f64 {
244 let device = match self.get_device(device_id) {
245 Some(d) => d,
246 None => return 0.0,
247 };
248
249 let base_throughput = match operation {
250 "add" | "subtract" | "multiply" => device.memory_bandwidth as f64 * 0.8,
251 "divide" | "sqrt" | "exp" | "log" => device.memory_bandwidth as f64 * 0.6,
252 "matrix_multiply" => (device.cores as f64 * device.clock_rate as f64 * 1e6) * 0.5,
253 "fft" => (device.cores as f64 * device.clock_rate as f64 * 1e6) * 0.3,
254 _ => device.memory_bandwidth as f64 * 0.5,
255 };
256
257 let array_factor = (array_size as f64).log2() / 20.0; base_throughput * (1.0 - array_factor.min(0.5))
259 }
260
261 pub fn should_use_gpu(&self, array_size: usize, operation: &str) -> bool {
263 if self.devices.is_empty() {
264 return false;
265 }
266
267 let min_size = match operation {
268 "add" | "subtract" | "multiply" | "divide" => 1000,
269 "matrix_multiply" => 100,
270 "fft" | "conv" => 512,
271 _ => 1000,
272 };
273
274 array_size >= min_size
275 }
276
277 pub fn get_utilization(&self) -> HashMap<u32, f64> {
279 let mut utilization = HashMap::new();
280
281 for device in &self.devices {
282 let recent_executions = self
284 .kernel_executions
285 .read()
286 .expect("operation should succeed")
287 .iter()
288 .filter(|e| e.device_id == device.id)
289 .filter(|e| e.execution_time > 0.0)
290 .count();
291
292 let util = (recent_executions as f64 / 10.0).min(1.0);
293 utilization.insert(device.id, util);
294 }
295
296 utilization
297 }
298
299 pub fn cleanup(&self) -> Result<(), GpuError> {
301 let allocations = self.allocations.read().expect("operation should succeed");
302 if !allocations.is_empty() {
303 return Err(GpuError::ResourcesNotFreed);
304 }
305
306 self.kernel_executions
308 .write()
309 .expect("operation should succeed")
310 .clear();
311 self.performance_counters
312 .write()
313 .expect("operation should succeed")
314 .clear();
315
316 Ok(())
317 }
318}
319
320#[derive(Debug, Clone)]
322pub struct GpuKernelInfo {
323 pub name: String,
324 pub device_id: u32,
325 pub grid_size: (u32, u32, u32),
326 pub block_size: (u32, u32, u32),
327 pub shared_memory: u32,
328 pub parameters: HashMap<String, String>,
329}
330
331#[derive(Debug, Clone)]
333pub struct MemoryStats {
334 pub total_memory: u64,
335 pub available_memory: u64,
336 pub allocated_memory: u64,
337 pub free_memory: u64,
338 pub num_allocations: usize,
339 pub largest_allocation: u64,
340 pub fragmentation_ratio: f64,
341}
342
343pub struct GpuArrayOps;
345
346impl GpuArrayOps {
347 pub fn add_arrays(a: &[f32], b: &[f32], _device_id: u32) -> Result<Vec<f32>, GpuError> {
349 if a.len() != b.len() {
350 return Err(GpuError::ShapeMismatch);
351 }
352
353 let result: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x + y).collect();
355 Ok(result)
356 }
357
358 pub fn multiply_arrays(a: &[f32], b: &[f32], _device_id: u32) -> Result<Vec<f32>, GpuError> {
360 if a.len() != b.len() {
361 return Err(GpuError::ShapeMismatch);
362 }
363
364 let result: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x * y).collect();
366 Ok(result)
367 }
368
369 pub fn matrix_multiply(
371 a: &[f32],
372 b: &[f32],
373 m: usize,
374 n: usize,
375 k: usize,
376 _device_id: u32,
377 ) -> Result<Vec<f32>, GpuError> {
378 if a.len() != m * k || b.len() != k * n {
379 return Err(GpuError::ShapeMismatch);
380 }
381
382 let mut result = vec![0.0f32; m * n];
384
385 for i in 0..m {
386 for j in 0..n {
387 for l in 0..k {
388 result[i * n + j] += a[i * k + l] * b[l * n + j];
389 }
390 }
391 }
392
393 Ok(result)
394 }
395
396 pub fn apply_activation(
398 input: &[f32],
399 activation: ActivationFunction,
400 _device_id: u32,
401 ) -> Result<Vec<f32>, GpuError> {
402 let result: Vec<f32> = input
404 .iter()
405 .map(|&x| {
406 match activation {
407 ActivationFunction::ReLU => x.max(0.0),
408 ActivationFunction::Sigmoid => 1.0 / (1.0 + (-x).exp()),
409 ActivationFunction::Tanh => x.tanh(),
410 ActivationFunction::Softmax => x.exp(), }
412 })
413 .collect();
414
415 Ok(result)
416 }
417
418 pub fn reduce_sum(input: &[f32], _device_id: u32) -> Result<f32, GpuError> {
420 Ok(input.iter().sum())
422 }
423
424 pub fn reduce_max(input: &[f32], _device_id: u32) -> Result<f32, GpuError> {
426 input
428 .iter()
429 .fold(f32::NEG_INFINITY, |a, &b| a.max(b))
430 .is_finite()
431 .then_some(input.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)))
432 .ok_or(GpuError::ComputationError)
433 }
434}
435
436#[derive(Debug, Clone, Copy)]
438pub enum ActivationFunction {
439 ReLU,
440 Sigmoid,
441 Tanh,
442 Softmax,
443}
444
445#[derive(Debug, thiserror::Error)]
447pub enum GpuError {
448 #[error("GPU device not found")]
449 DeviceNotFound,
450 #[error("Out of GPU memory")]
451 OutOfMemory,
452 #[error("Invalid GPU pointer")]
453 InvalidPointer,
454 #[error("GPU computation error")]
455 ComputationError,
456 #[error("Array shape mismatch")]
457 ShapeMismatch,
458 #[error("GPU resources not freed")]
459 ResourcesNotFreed,
460 #[error("GPU initialization failed: {0}")]
461 InitializationFailed(String),
462}
463
464#[derive(Debug)]
466pub struct GpuProfiler {
467 kernel_times: HashMap<String, Vec<f64>>,
468 memory_transfers: Vec<(Instant, u64, String)>,
469 device_utilization: HashMap<u32, Vec<(Instant, f64)>>,
470}
471
472impl GpuProfiler {
473 pub fn new() -> Self {
475 Self {
476 kernel_times: HashMap::new(),
477 memory_transfers: Vec::new(),
478 device_utilization: HashMap::new(),
479 }
480 }
481
482 pub fn record_kernel_time(&mut self, kernel_name: &str, time_ms: f64) {
484 self.kernel_times
485 .entry(kernel_name.to_string())
486 .or_default()
487 .push(time_ms);
488 }
489
490 pub fn record_memory_transfer(&mut self, size: u64, direction: &str) {
492 self.memory_transfers
493 .push((Instant::now(), size, direction.to_string()));
494 }
495
496 pub fn record_utilization(&mut self, device_id: u32, utilization: f64) {
498 self.device_utilization
499 .entry(device_id)
500 .or_default()
501 .push((Instant::now(), utilization));
502 }
503
504 pub fn get_kernel_stats(&self) -> HashMap<String, KernelStats> {
506 let mut stats = HashMap::new();
507
508 for (kernel_name, times) in &self.kernel_times {
509 let count = times.len();
510 let total_time: f64 = times.iter().sum();
511 let avg_time = total_time / count as f64;
512 let min_time = times.iter().fold(f64::INFINITY, |a, &b| a.min(b));
513 let max_time = times.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
514
515 stats.insert(
516 kernel_name.clone(),
517 KernelStats {
518 count,
519 total_time,
520 avg_time,
521 min_time,
522 max_time,
523 },
524 );
525 }
526
527 stats
528 }
529
530 pub fn get_memory_transfer_stats(&self) -> MemoryTransferStats {
532 let total_transfers = self.memory_transfers.len();
533 let total_bytes: u64 = self.memory_transfers.iter().map(|(_, size, _)| size).sum();
534
535 let host_to_device = self
536 .memory_transfers
537 .iter()
538 .filter(|(_, _, dir)| dir == "host_to_device")
539 .count();
540
541 let device_to_host = self
542 .memory_transfers
543 .iter()
544 .filter(|(_, _, dir)| dir == "device_to_host")
545 .count();
546
547 MemoryTransferStats {
548 total_transfers,
549 total_bytes,
550 host_to_device_transfers: host_to_device,
551 device_to_host_transfers: device_to_host,
552 }
553 }
554
555 pub fn clear(&mut self) {
557 self.kernel_times.clear();
558 self.memory_transfers.clear();
559 self.device_utilization.clear();
560 }
561}
562
563#[derive(Debug, Clone)]
565pub struct KernelStats {
566 pub count: usize,
567 pub total_time: f64,
568 pub avg_time: f64,
569 pub min_time: f64,
570 pub max_time: f64,
571}
572
573#[derive(Debug, Clone)]
575pub struct MemoryTransferStats {
576 pub total_transfers: usize,
577 pub total_bytes: u64,
578 pub host_to_device_transfers: usize,
579 pub device_to_host_transfers: usize,
580}
581
582impl Default for GpuUtils {
583 fn default() -> Self {
584 Self::new()
585 }
586}
587
588impl Default for GpuProfiler {
589 fn default() -> Self {
590 Self::new()
591 }
592}
593
594pub struct MultiGpuCoordinator {
596 gpus: HashMap<u32, GpuUtils>,
597 load_balancer: LoadBalancer,
598 #[allow(dead_code)]
599 communication_topology: CommunicationTopology,
600 #[allow(dead_code)]
601 synchronization_barriers: Vec<SynchronizationBarrier>,
602}
603
604impl Default for MultiGpuCoordinator {
605 fn default() -> Self {
606 Self::new()
607 }
608}
609
610impl MultiGpuCoordinator {
611 pub fn new() -> Self {
613 Self {
614 gpus: HashMap::new(),
615 load_balancer: LoadBalancer::new(),
616 communication_topology: CommunicationTopology::Ring,
617 synchronization_barriers: Vec::new(),
618 }
619 }
620
621 pub fn init_all_gpus(&mut self) -> Result<(), GpuError> {
623 for gpu_id in 0..8 {
624 let mut gpu = GpuUtils::new();
626 if gpu.init_devices().is_ok() && !gpu.devices.is_empty() {
627 self.gpus.insert(gpu_id, gpu);
628 }
629 }
630
631 if self.gpus.is_empty() {
632 return Err(GpuError::InitializationFailed("No GPUs found".to_string()));
633 }
634
635 Ok(())
636 }
637
638 pub fn get_optimal_assignment(&self, workload: &DistributedWorkload) -> Vec<GpuAssignment> {
640 self.load_balancer.assign_workload(workload, &self.gpus)
641 }
642
643 pub fn execute_distributed(
645 &self,
646 operation: &DistributedOperation,
647 ) -> Result<DistributedResult, GpuError> {
648 let assignments = self.get_optimal_assignment(&operation.workload);
649 let mut results = Vec::new();
650
651 for assignment in assignments {
653 let gpu = self
654 .gpus
655 .get(&assignment.gpu_id)
656 .ok_or(GpuError::DeviceNotFound)?;
657
658 let kernel_info = GpuKernelInfo {
659 name: operation.kernel_name.clone(),
660 device_id: assignment.gpu_id,
661 grid_size: assignment.grid_size,
662 block_size: assignment.block_size,
663 shared_memory: assignment.shared_memory,
664 parameters: assignment.parameters.clone(),
665 };
666
667 let execution = gpu.execute_kernel(&kernel_info)?;
668 results.push(execution);
669 }
670
671 let total_time: f64 = results.iter().map(|e| e.execution_time).sum();
673 Ok(DistributedResult {
674 executions: results,
675 total_time,
676 communication_overhead: 0.0, })
678 }
679
680 pub fn synchronize_all(&self) -> Result<(), GpuError> {
682 std::thread::sleep(std::time::Duration::from_millis(1));
684 Ok(())
685 }
686
687 pub fn get_cluster_memory_stats(&self) -> ClusterMemoryStats {
689 let mut total_memory = 0;
690 let mut total_allocated = 0;
691 let mut device_stats = HashMap::new();
692
693 for (gpu_id, gpu) in &self.gpus {
694 let stats = gpu.get_memory_stats();
695 if let Some(stat) = stats.get(gpu_id) {
696 total_memory += stat.total_memory;
697 total_allocated += stat.allocated_memory;
698 device_stats.insert(*gpu_id, stat.clone());
699 }
700 }
701
702 ClusterMemoryStats {
703 total_memory,
704 total_allocated,
705 total_free: total_memory - total_allocated,
706 num_devices: self.gpus.len(),
707 device_stats,
708 }
709 }
710}
711
712pub struct GpuMemoryPool {
714 pools: HashMap<u32, Vec<MemoryBlock>>,
715 #[allow(dead_code)]
716 allocation_strategy: AllocationStrategy,
717 #[allow(dead_code)]
718 fragmentation_threshold: f64,
719}
720
721impl GpuMemoryPool {
722 pub fn new(strategy: AllocationStrategy) -> Self {
724 Self {
725 pools: HashMap::new(),
726 allocation_strategy: strategy,
727 fragmentation_threshold: 0.3,
728 }
729 }
730
731 pub fn allocate(&mut self, size: u64, device_id: u32) -> Result<u64, GpuError> {
733 let pool = self.pools.entry(device_id).or_default();
735
736 for (i, block) in pool.iter().enumerate() {
738 if !block.is_allocated && block.size >= size {
739 if block.size > size * 2 {
741 let new_block = MemoryBlock {
742 ptr: block.ptr + size,
743 size: block.size - size,
744 is_allocated: false,
745 allocation_time: None,
746 };
747 pool.push(new_block);
748
749 pool[i].size = size;
750 }
751
752 pool[i].is_allocated = true;
753 pool[i].allocation_time = Some(Instant::now());
754 return Ok(pool[i].ptr);
755 }
756 }
757
758 let ptr = self.allocate_new_block(size, device_id)?;
760
761 let pool = self.pools.entry(device_id).or_default();
763 pool.push(MemoryBlock {
764 ptr,
765 size,
766 is_allocated: true,
767 allocation_time: Some(Instant::now()),
768 });
769
770 Ok(ptr)
771 }
772
773 pub fn free(&mut self, ptr: u64, device_id: u32) -> Result<(), GpuError> {
775 let pool = self
776 .pools
777 .get_mut(&device_id)
778 .ok_or(GpuError::DeviceNotFound)?;
779
780 for block in pool.iter_mut() {
781 if block.ptr == ptr {
782 block.is_allocated = false;
783 block.allocation_time = None;
784 self.try_merge_blocks(device_id);
785 return Ok(());
786 }
787 }
788
789 Err(GpuError::InvalidPointer)
790 }
791
792 pub fn defragment(&mut self, device_id: u32) -> Result<DefragmentationResult, GpuError> {
794 let before_fragmentation = self.calculate_fragmentation(device_id);
795
796 let pool = self
797 .pools
798 .get_mut(&device_id)
799 .ok_or(GpuError::DeviceNotFound)?;
800 let before_blocks = pool.len();
801
802 pool.sort_by_key(|b| b.ptr);
804
805 let mut i = 0;
807 while i < pool.len() - 1 {
808 if !pool[i].is_allocated
809 && !pool[i + 1].is_allocated
810 && pool[i].ptr + pool[i].size == pool[i + 1].ptr
811 {
812 pool[i].size += pool[i + 1].size;
813 pool.remove(i + 1);
814 } else {
815 i += 1;
816 }
817 }
818
819 let after_blocks = pool.len();
820 let after_fragmentation = self.calculate_fragmentation(device_id);
821
822 Ok(DefragmentationResult {
823 blocks_before: before_blocks,
824 blocks_after: after_blocks,
825 fragmentation_before: before_fragmentation,
826 fragmentation_after: after_fragmentation,
827 })
828 }
829
830 fn allocate_new_block(&self, size: u64, _device_id: u32) -> Result<u64, GpuError> {
831 let ptr = (std::ptr::null::<u8>() as u64) + size;
833 Ok(ptr)
834 }
835
836 fn try_merge_blocks(&mut self, device_id: u32) {
837 if let Some(pool) = self.pools.get_mut(&device_id) {
838 pool.sort_by_key(|b| b.ptr);
839
840 let mut i = 0;
841 while i < pool.len() - 1 {
842 if !pool[i].is_allocated
843 && !pool[i + 1].is_allocated
844 && pool[i].ptr + pool[i].size == pool[i + 1].ptr
845 {
846 pool[i].size += pool[i + 1].size;
847 pool.remove(i + 1);
848 } else {
849 i += 1;
850 }
851 }
852 }
853 }
854
855 fn calculate_fragmentation(&self, device_id: u32) -> f64 {
856 let empty_pool = Vec::new();
857 let pool = self.pools.get(&device_id).unwrap_or(&empty_pool);
858 let free_blocks = pool.iter().filter(|b| !b.is_allocated).count();
859 let total_blocks = pool.len();
860
861 if total_blocks == 0 {
862 0.0
863 } else {
864 free_blocks as f64 / total_blocks as f64
865 }
866 }
867}
868
869pub struct AsyncGpuOps {
871 streams: HashMap<u32, Vec<GpuStream>>,
872 pending_operations: Vec<AsyncOperation>,
873}
874
875impl Default for AsyncGpuOps {
876 fn default() -> Self {
877 Self::new()
878 }
879}
880
881impl AsyncGpuOps {
882 pub fn new() -> Self {
884 Self {
885 streams: HashMap::new(),
886 pending_operations: Vec::new(),
887 }
888 }
889
890 pub fn create_stream(&mut self, device_id: u32) -> Result<u32, GpuError> {
892 let stream_id = self.streams.get(&device_id).map_or(0, |s| s.len() as u32);
893 let stream = GpuStream {
894 id: stream_id,
895 device_id,
896 is_busy: false,
897 priority: StreamPriority::Normal,
898 };
899
900 self.streams.entry(device_id).or_default().push(stream);
901 Ok(stream_id)
902 }
903
904 pub fn launch_kernel_async(
906 &mut self,
907 kernel: &GpuKernelInfo,
908 stream_id: u32,
909 ) -> Result<AsyncOperationHandle, GpuError> {
910 let operation = AsyncOperation {
911 id: self.pending_operations.len() as u32,
912 kernel_info: kernel.clone(),
913 stream_id,
914 start_time: Instant::now(),
915 status: OperationStatus::Pending,
916 };
917
918 let handle = AsyncOperationHandle {
919 operation_id: operation.id,
920 device_id: kernel.device_id,
921 };
922
923 self.pending_operations.push(operation);
924 Ok(handle)
925 }
926
927 pub fn wait_for_completion(
929 &mut self,
930 handle: &AsyncOperationHandle,
931 ) -> Result<GpuKernelExecution, GpuError> {
932 std::thread::sleep(std::time::Duration::from_millis(1));
934
935 if let Some(op) = self
936 .pending_operations
937 .iter_mut()
938 .find(|op| op.id == handle.operation_id)
939 {
940 op.status = OperationStatus::Completed;
941
942 Ok(GpuKernelExecution {
943 kernel_name: op.kernel_info.name.clone(),
944 device_id: op.kernel_info.device_id,
945 grid_size: op.kernel_info.grid_size,
946 block_size: op.kernel_info.block_size,
947 shared_memory: op.kernel_info.shared_memory,
948 execution_time: op.start_time.elapsed().as_secs_f64() * 1000.0,
949 parameters: op.kernel_info.parameters.clone(),
950 })
951 } else {
952 Err(GpuError::ComputationError)
953 }
954 }
955
956 pub fn is_complete(&self, handle: &AsyncOperationHandle) -> bool {
958 self.pending_operations
959 .iter()
960 .find(|op| op.id == handle.operation_id)
961 .is_some_and(|op| matches!(op.status, OperationStatus::Completed))
962 }
963}
964
965pub struct GpuOptimizationAdvisor {
967 performance_history: HashMap<String, Vec<PerformanceMetric>>,
968 optimization_rules: Vec<OptimizationRule>,
969}
970
971impl Default for GpuOptimizationAdvisor {
972 fn default() -> Self {
973 Self::new()
974 }
975}
976
977impl GpuOptimizationAdvisor {
978 pub fn new() -> Self {
980 let mut advisor = Self {
981 performance_history: HashMap::new(),
982 optimization_rules: Vec::new(),
983 };
984
985 advisor.init_default_rules();
986 advisor
987 }
988
989 pub fn analyze_performance(
991 &mut self,
992 kernel_name: &str,
993 execution: &GpuKernelExecution,
994 workload_size: usize,
995 ) -> Vec<OptimizationRecommendation> {
996 let metric = PerformanceMetric {
997 execution_time: execution.execution_time,
998 throughput: workload_size as f64 / execution.execution_time,
999 memory_bandwidth: 0.0, occupancy: self.calculate_occupancy(execution),
1001 };
1002
1003 self.performance_history
1004 .entry(kernel_name.to_string())
1005 .or_default()
1006 .push(metric.clone());
1007
1008 let mut recommendations = Vec::new();
1009
1010 for rule in &self.optimization_rules {
1011 if let Some(recommendation) = rule.evaluate(&metric, execution) {
1012 recommendations.push(recommendation);
1013 }
1014 }
1015
1016 recommendations
1017 }
1018
1019 fn init_default_rules(&mut self) {
1020 self.optimization_rules.push(OptimizationRule {
1021 name: "Low Occupancy".to_string(),
1022 condition: Box::new(|metric, _| metric.occupancy < 0.5),
1023 recommendation: "Consider increasing block size or reducing register usage".to_string(),
1024 priority: RecommendationPriority::High,
1025 });
1026
1027 self.optimization_rules.push(OptimizationRule {
1028 name: "Memory Bandwidth".to_string(),
1029 condition: Box::new(|metric, _| metric.memory_bandwidth < 0.7),
1030 recommendation: "Optimize memory access patterns for better coalescing".to_string(),
1031 priority: RecommendationPriority::Medium,
1032 });
1033
1034 self.optimization_rules.push(OptimizationRule {
1035 name: "Small Grid Size".to_string(),
1036 condition: Box::new(|_, execution| {
1037 let total_threads = execution.grid_size.0
1038 * execution.grid_size.1
1039 * execution.grid_size.2
1040 * execution.block_size.0
1041 * execution.block_size.1
1042 * execution.block_size.2;
1043 total_threads < 1024
1044 }),
1045 recommendation: "Consider increasing grid size to better utilize GPU cores".to_string(),
1046 priority: RecommendationPriority::Low,
1047 });
1048 }
1049
1050 fn calculate_occupancy(&self, execution: &GpuKernelExecution) -> f64 {
1051 let threads_per_block =
1052 execution.block_size.0 * execution.block_size.1 * execution.block_size.2;
1053 let blocks_per_sm = 2048 / threads_per_block.max(1); (blocks_per_sm as f64 / 32.0).min(1.0) }
1056}
1057
1058#[derive(Debug, Clone)]
1061pub struct DistributedWorkload {
1062 pub total_elements: usize,
1063 pub operation_type: String,
1064 pub memory_requirement: u64,
1065 pub computation_complexity: f64,
1066}
1067
1068#[derive(Debug, Clone)]
1069pub struct DistributedOperation {
1070 pub kernel_name: String,
1071 pub workload: DistributedWorkload,
1072}
1073
1074#[derive(Debug, Clone)]
1075pub struct DistributedResult {
1076 pub executions: Vec<GpuKernelExecution>,
1077 pub total_time: f64,
1078 pub communication_overhead: f64,
1079}
1080
1081#[derive(Debug, Clone)]
1082pub struct GpuAssignment {
1083 pub gpu_id: u32,
1084 pub grid_size: (u32, u32, u32),
1085 pub block_size: (u32, u32, u32),
1086 pub shared_memory: u32,
1087 pub parameters: HashMap<String, String>,
1088}
1089
1090#[derive(Debug, Clone)]
1091pub struct LoadBalancer {
1092 #[allow(dead_code)]
1093 strategy: LoadBalancingStrategy,
1094}
1095
1096impl Default for LoadBalancer {
1097 fn default() -> Self {
1098 Self::new()
1099 }
1100}
1101
1102impl LoadBalancer {
1103 pub fn new() -> Self {
1104 Self {
1105 strategy: LoadBalancingStrategy::WorkloadProportional,
1106 }
1107 }
1108
1109 pub fn assign_workload(
1110 &self,
1111 workload: &DistributedWorkload,
1112 gpus: &HashMap<u32, GpuUtils>,
1113 ) -> Vec<GpuAssignment> {
1114 let mut assignments = Vec::new();
1115 let num_gpus = gpus.len() as u32;
1116
1117 if num_gpus == 0 {
1118 return assignments;
1119 }
1120
1121 let elements_per_gpu = workload.total_elements / num_gpus as usize;
1122
1123 for (gpu_id, _) in gpus.iter() {
1124 let assignment = GpuAssignment {
1125 gpu_id: *gpu_id,
1126 grid_size: (elements_per_gpu as u32 / 256, 1, 1),
1127 block_size: (256, 1, 1),
1128 shared_memory: 0,
1129 parameters: HashMap::new(),
1130 };
1131 assignments.push(assignment);
1132 }
1133
1134 assignments
1135 }
1136}
1137
1138#[derive(Debug, Clone)]
1139pub enum LoadBalancingStrategy {
1140 RoundRobin,
1141 WorkloadProportional,
1142 MemoryAware,
1143 PerformanceBased,
1144}
1145
1146#[derive(Debug, Clone)]
1147pub enum CommunicationTopology {
1148 Ring,
1149 Tree,
1150 AllToAll,
1151 Custom(Vec<Vec<u32>>),
1152}
1153
1154#[derive(Debug, Clone)]
1155pub struct SynchronizationBarrier {
1156 pub id: u32,
1157 pub participating_gpus: Vec<u32>,
1158 pub barrier_type: BarrierType,
1159}
1160
1161#[derive(Debug, Clone)]
1162pub enum BarrierType {
1163 Global,
1164 Local(Vec<u32>),
1165 Hierarchical,
1166}
1167
1168#[derive(Debug, Clone)]
1169pub struct ClusterMemoryStats {
1170 pub total_memory: u64,
1171 pub total_allocated: u64,
1172 pub total_free: u64,
1173 pub num_devices: usize,
1174 pub device_stats: HashMap<u32, MemoryStats>,
1175}
1176
1177#[derive(Debug, Clone)]
1178pub struct MemoryBlock {
1179 pub ptr: u64,
1180 pub size: u64,
1181 pub is_allocated: bool,
1182 pub allocation_time: Option<Instant>,
1183}
1184
1185#[derive(Debug, Clone)]
1186pub enum AllocationStrategy {
1187 FirstFit,
1188 BestFit,
1189 WorstFit,
1190 BuddySystem,
1191}
1192
1193#[derive(Debug, Clone)]
1194pub struct DefragmentationResult {
1195 pub blocks_before: usize,
1196 pub blocks_after: usize,
1197 pub fragmentation_before: f64,
1198 pub fragmentation_after: f64,
1199}
1200
1201#[derive(Debug, Clone)]
1202pub struct GpuStream {
1203 pub id: u32,
1204 pub device_id: u32,
1205 pub is_busy: bool,
1206 pub priority: StreamPriority,
1207}
1208
1209#[derive(Debug, Clone)]
1210pub enum StreamPriority {
1211 Low,
1212 Normal,
1213 High,
1214}
1215
1216#[derive(Debug, Clone)]
1217pub struct AsyncOperation {
1218 pub id: u32,
1219 pub kernel_info: GpuKernelInfo,
1220 pub stream_id: u32,
1221 pub start_time: Instant,
1222 pub status: OperationStatus,
1223}
1224
1225#[derive(Debug, Clone)]
1226pub enum OperationStatus {
1227 Pending,
1228 Running,
1229 Completed,
1230 Failed,
1231}
1232
1233#[derive(Debug, Clone)]
1234pub struct AsyncOperationHandle {
1235 pub operation_id: u32,
1236 pub device_id: u32,
1237}
1238
1239#[derive(Debug, Clone)]
1240pub struct PerformanceMetric {
1241 pub execution_time: f64,
1242 pub throughput: f64,
1243 pub memory_bandwidth: f64,
1244 pub occupancy: f64,
1245}
1246
1247type OptimizationCondition =
1248 Box<dyn Fn(&PerformanceMetric, &GpuKernelExecution) -> bool + Send + Sync>;
1249
1250pub struct OptimizationRule {
1251 pub name: String,
1252 pub condition: OptimizationCondition,
1253 pub recommendation: String,
1254 pub priority: RecommendationPriority,
1255}
1256
1257impl std::fmt::Debug for OptimizationRule {
1258 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1259 f.debug_struct("OptimizationRule")
1260 .field("name", &self.name)
1261 .field("condition", &"<function>")
1262 .field("recommendation", &self.recommendation)
1263 .field("priority", &self.priority)
1264 .finish()
1265 }
1266}
1267
1268impl Clone for OptimizationRule {
1269 fn clone(&self) -> Self {
1270 OptimizationRule {
1274 name: self.name.clone(),
1275 condition: Box::new(|_metric, _execution| false), recommendation: self.recommendation.clone(),
1277 priority: self.priority.clone(),
1278 }
1279 }
1280}
1281
1282impl OptimizationRule {
1283 pub fn evaluate(
1284 &self,
1285 metric: &PerformanceMetric,
1286 execution: &GpuKernelExecution,
1287 ) -> Option<OptimizationRecommendation> {
1288 if (self.condition)(metric, execution) {
1289 Some(OptimizationRecommendation {
1290 rule_name: self.name.clone(),
1291 recommendation: self.recommendation.clone(),
1292 priority: self.priority.clone(),
1293 estimated_improvement: 0.0, })
1295 } else {
1296 None
1297 }
1298 }
1299}
1300
1301#[derive(Debug, Clone)]
1302pub struct OptimizationRecommendation {
1303 pub rule_name: String,
1304 pub recommendation: String,
1305 pub priority: RecommendationPriority,
1306 pub estimated_improvement: f64,
1307}
1308
1309#[derive(Debug, Clone)]
1310pub enum RecommendationPriority {
1311 Low,
1312 Medium,
1313 High,
1314 Critical,
1315}
1316
1317#[allow(non_snake_case)]
1318#[cfg(test)]
1319mod tests {
1320 use super::*;
1321
1322 #[test]
1323 fn test_gpu_utils_creation() {
1324 let utils = GpuUtils::new();
1325 assert!(utils.devices.is_empty());
1326 assert!(utils
1327 .allocations
1328 .read()
1329 .expect("operation should succeed")
1330 .is_empty());
1331 }
1332
1333 #[test]
1334 fn test_device_initialization() {
1335 let mut utils = GpuUtils::new();
1336 assert!(utils.init_devices().is_ok());
1337 assert!(!utils.devices.is_empty());
1338 }
1339
1340 #[test]
1341 fn test_device_selection() {
1342 let mut utils = GpuUtils::new();
1343 utils.init_devices().expect("operation should succeed");
1344
1345 let best_device = utils.get_best_device();
1346 assert!(best_device.is_some());
1347 assert!(!best_device.expect("operation should succeed").is_integrated);
1348 }
1349
1350 #[test]
1351 fn test_memory_allocation() {
1352 let mut utils = GpuUtils::new();
1353 utils.init_devices().expect("operation should succeed");
1354
1355 let ptr = utils
1356 .allocate_memory(1024, 0, "test")
1357 .expect("operation should succeed");
1358 assert!(ptr > 0);
1359
1360 assert!(utils.free_memory(ptr).is_ok());
1361 }
1362
1363 #[test]
1364 fn test_kernel_execution() {
1365 let mut utils = GpuUtils::new();
1366 utils.init_devices().expect("operation should succeed");
1367
1368 let kernel_info = GpuKernelInfo {
1369 name: "test_kernel".to_string(),
1370 device_id: 0,
1371 grid_size: (1, 1, 1),
1372 block_size: (256, 1, 1),
1373 shared_memory: 0,
1374 parameters: HashMap::new(),
1375 };
1376
1377 let execution = utils
1378 .execute_kernel(&kernel_info)
1379 .expect("operation should succeed");
1380 assert_eq!(execution.kernel_name, "test_kernel");
1381 assert!(execution.execution_time > 0.0);
1382 }
1383
1384 #[test]
1385 fn test_array_operations() {
1386 let a = vec![1.0, 2.0, 3.0, 4.0];
1387 let b = vec![5.0, 6.0, 7.0, 8.0];
1388
1389 let result = GpuArrayOps::add_arrays(&a, &b, 0).expect("operation should succeed");
1390 assert_eq!(result, vec![6.0, 8.0, 10.0, 12.0]);
1391
1392 let result = GpuArrayOps::multiply_arrays(&a, &b, 0).expect("operation should succeed");
1393 assert_eq!(result, vec![5.0, 12.0, 21.0, 32.0]);
1394 }
1395
1396 #[test]
1397 fn test_matrix_multiplication() {
1398 let a = vec![1.0, 2.0, 3.0, 4.0]; let b = vec![5.0, 6.0, 7.0, 8.0]; let result =
1402 GpuArrayOps::matrix_multiply(&a, &b, 2, 2, 2, 0).expect("operation should succeed");
1403 assert_eq!(result, vec![19.0, 22.0, 43.0, 50.0]);
1404 }
1405
1406 #[test]
1407 fn test_activation_functions() {
1408 let input = vec![-1.0, 0.0, 1.0, 2.0];
1409
1410 let result = GpuArrayOps::apply_activation(&input, ActivationFunction::ReLU, 0)
1411 .expect("operation should succeed");
1412 assert_eq!(result, vec![0.0, 0.0, 1.0, 2.0]);
1413
1414 let result = GpuArrayOps::apply_activation(&input, ActivationFunction::Sigmoid, 0)
1415 .expect("operation should succeed");
1416 assert!(result.iter().all(|&x| (0.0..=1.0).contains(&x)));
1417 }
1418
1419 #[test]
1420 fn test_reduction_operations() {
1421 let input = vec![1.0, 2.0, 3.0, 4.0, 5.0];
1422
1423 let sum = GpuArrayOps::reduce_sum(&input, 0).expect("operation should succeed");
1424 assert_eq!(sum, 15.0);
1425
1426 let max = GpuArrayOps::reduce_max(&input, 0).expect("operation should succeed");
1427 assert_eq!(max, 5.0);
1428 }
1429
1430 #[test]
1431 fn test_gpu_profiler() {
1432 let mut profiler = GpuProfiler::new();
1433
1434 profiler.record_kernel_time("test_kernel", 1.5);
1435 profiler.record_kernel_time("test_kernel", 2.0);
1436 profiler.record_memory_transfer(1024, "host_to_device");
1437
1438 let stats = profiler.get_kernel_stats();
1439 assert!(stats.contains_key("test_kernel"));
1440 assert_eq!(stats["test_kernel"].count, 2);
1441 assert_eq!(stats["test_kernel"].avg_time, 1.75);
1442
1443 let mem_stats = profiler.get_memory_transfer_stats();
1444 assert_eq!(mem_stats.total_transfers, 1);
1445 assert_eq!(mem_stats.total_bytes, 1024);
1446 }
1447
1448 #[test]
1449 fn test_throughput_estimation() {
1450 let mut utils = GpuUtils::new();
1451 utils.init_devices().expect("operation should succeed");
1452
1453 let throughput = utils.estimate_throughput(0, 1000, "add");
1454 assert!(throughput > 0.0);
1455
1456 let should_use = utils.should_use_gpu(1000, "add");
1457 assert!(should_use);
1458
1459 let should_not_use = utils.should_use_gpu(100, "add");
1460 assert!(!should_not_use);
1461 }
1462
1463 #[test]
1464 fn test_memory_stats() {
1465 let mut utils = GpuUtils::new();
1466 utils.init_devices().expect("operation should succeed");
1467
1468 let _ptr = utils
1469 .allocate_memory(1024, 0, "test")
1470 .expect("operation should succeed");
1471 let stats = utils.get_memory_stats();
1472
1473 assert!(stats.contains_key(&0));
1474 assert_eq!(stats[&0].allocated_memory, 1024);
1475 assert_eq!(stats[&0].num_allocations, 1);
1476 }
1477
1478 #[test]
1479 fn test_error_handling() {
1480 let utils = GpuUtils::new();
1481
1482 let result = utils.allocate_memory(1024, 999, "test");
1484 assert!(matches!(result, Err(GpuError::DeviceNotFound)));
1485
1486 let result = utils.free_memory(0);
1488 assert!(matches!(result, Err(GpuError::InvalidPointer)));
1489
1490 let a = vec![1.0, 2.0];
1492 let b = vec![3.0, 4.0, 5.0];
1493 let result = GpuArrayOps::add_arrays(&a, &b, 0);
1494 assert!(matches!(result, Err(GpuError::ShapeMismatch)));
1495 }
1496
1497 #[test]
1500 fn test_multi_gpu_coordinator() {
1501 let mut coordinator = MultiGpuCoordinator::new();
1502
1503 let result = coordinator.init_all_gpus();
1505 assert!(result.is_ok() || matches!(result, Err(GpuError::InitializationFailed(_))));
1506
1507 let workload = DistributedWorkload {
1509 total_elements: 10_000,
1510 operation_type: "matrix_multiply".to_string(),
1511 memory_requirement: 1024 * 1024,
1512 computation_complexity: 1.0,
1513 };
1514
1515 let assignments = coordinator.get_optimal_assignment(&workload);
1516 assert!(!assignments.is_empty() || coordinator.gpus.is_empty());
1517 }
1518
1519 #[test]
1520 fn test_distributed_operation() {
1521 let mut coordinator = MultiGpuCoordinator::new();
1522 let init_result = coordinator.init_all_gpus();
1523
1524 let operation = DistributedOperation {
1525 kernel_name: "test_kernel".to_string(),
1526 workload: DistributedWorkload {
1527 total_elements: 1000,
1528 operation_type: "add".to_string(),
1529 memory_requirement: 4000,
1530 computation_complexity: 0.5,
1531 },
1532 };
1533
1534 if init_result.is_ok() && !coordinator.gpus.is_empty() {
1535 let result = coordinator.execute_distributed(&operation);
1536
1537 if let Ok(dist_result) = result {
1540 assert!(!dist_result.executions.is_empty());
1541 assert!(dist_result.total_time >= 0.0);
1542 } else {
1543 assert!(!coordinator.gpus.is_empty());
1546 }
1547 } else {
1548 assert!(coordinator.gpus.is_empty());
1551 }
1552 }
1553
1554 #[test]
1555 fn test_cluster_memory_stats() {
1556 let mut coordinator = MultiGpuCoordinator::new();
1557 let _ = coordinator.init_all_gpus();
1558
1559 let stats = coordinator.get_cluster_memory_stats();
1560 assert_eq!(stats.num_devices, coordinator.gpus.len());
1561 assert_eq!(stats.total_free, stats.total_memory - stats.total_allocated);
1562 }
1563
1564 #[test]
1565 fn test_gpu_memory_pool() {
1566 let mut pool = GpuMemoryPool::new(AllocationStrategy::FirstFit);
1567
1568 let ptr1 = pool.allocate(1024, 0);
1570 assert!(ptr1.is_ok());
1571
1572 let ptr2 = pool.allocate(2048, 0);
1573 assert!(ptr2.is_ok());
1574
1575 let free_result = pool.free(ptr1.expect("operation should succeed"), 0);
1577 assert!(free_result.is_ok());
1578
1579 let defrag_result = pool.defragment(0);
1581 assert!(defrag_result.is_ok());
1582
1583 let defrag = defrag_result.expect("operation should succeed");
1584 assert!(defrag.fragmentation_after <= defrag.fragmentation_before);
1585 }
1586
1587 #[test]
1588 fn test_memory_pool_strategies() {
1589 let strategies = vec![
1590 AllocationStrategy::FirstFit,
1591 AllocationStrategy::BestFit,
1592 AllocationStrategy::WorstFit,
1593 AllocationStrategy::BuddySystem,
1594 ];
1595
1596 for strategy in strategies {
1597 let mut pool = GpuMemoryPool::new(strategy);
1598 let ptr = pool.allocate(1024, 0);
1599 assert!(ptr.is_ok());
1600 }
1601 }
1602
1603 #[test]
1604 fn test_async_gpu_operations() {
1605 let mut async_ops = AsyncGpuOps::new();
1606
1607 let stream_id = async_ops.create_stream(0);
1609 assert!(stream_id.is_ok());
1610
1611 let kernel_info = GpuKernelInfo {
1613 name: "async_test".to_string(),
1614 device_id: 0,
1615 grid_size: (1, 1, 1),
1616 block_size: (256, 1, 1),
1617 shared_memory: 0,
1618 parameters: HashMap::new(),
1619 };
1620
1621 let handle = async_ops
1622 .launch_kernel_async(&kernel_info, stream_id.expect("operation should succeed"));
1623 assert!(handle.is_ok());
1624
1625 let operation_handle = handle.expect("operation should succeed");
1626
1627 let _is_complete_before = async_ops.is_complete(&operation_handle);
1629
1630 let execution = async_ops.wait_for_completion(&operation_handle);
1632 assert!(execution.is_ok());
1633
1634 let is_complete_after = async_ops.is_complete(&operation_handle);
1635 assert!(is_complete_after);
1636 }
1637
1638 #[test]
1639 fn test_gpu_optimization_advisor() {
1640 let mut advisor = GpuOptimizationAdvisor::new();
1641
1642 let execution = GpuKernelExecution {
1644 kernel_name: "test_kernel".to_string(),
1645 device_id: 0,
1646 grid_size: (10, 1, 1), block_size: (32, 1, 1), shared_memory: 0,
1649 execution_time: 5.0,
1650 parameters: HashMap::new(),
1651 };
1652
1653 let recommendations = advisor.analyze_performance("test_kernel", &execution, 1000);
1654 assert!(!recommendations.is_empty());
1655
1656 let has_grid_size_recommendation = recommendations
1658 .iter()
1659 .any(|r| r.rule_name.contains("Grid Size"));
1660 assert!(has_grid_size_recommendation);
1661 }
1662
1663 #[test]
1664 fn test_load_balancer() {
1665 let balancer = LoadBalancer::new();
1666 let mut gpus = HashMap::new();
1667
1668 let mut gpu1 = GpuUtils::new();
1670 let mut gpu2 = GpuUtils::new();
1671 let _ = gpu1.init_devices();
1672 let _ = gpu2.init_devices();
1673
1674 gpus.insert(0, gpu1);
1675 gpus.insert(1, gpu2);
1676
1677 let workload = DistributedWorkload {
1678 total_elements: 10_000,
1679 operation_type: "matrix_multiply".to_string(),
1680 memory_requirement: 1024 * 1024,
1681 computation_complexity: 1.0,
1682 };
1683
1684 let assignments = balancer.assign_workload(&workload, &gpus);
1685 assert_eq!(assignments.len(), gpus.len());
1686
1687 let total_elements: u32 = assignments
1689 .iter()
1690 .map(|a| a.grid_size.0 * a.block_size.0)
1691 .sum();
1692 assert!(total_elements > 0);
1693 }
1694
1695 #[test]
1696 fn test_stream_priorities() {
1697 let mut async_ops = AsyncGpuOps::new();
1698 let _stream_id = async_ops
1699 .create_stream(0)
1700 .expect("operation should succeed");
1701
1702 let streams = async_ops.streams.get(&0).expect("operation should succeed");
1704 assert_eq!(streams.len(), 1);
1705 assert!(matches!(streams[0].priority, StreamPriority::Normal));
1706 }
1707
1708 #[test]
1709 fn test_memory_block_operations() {
1710 let block1 = MemoryBlock {
1711 ptr: 1000,
1712 size: 1024,
1713 is_allocated: false,
1714 allocation_time: None,
1715 };
1716
1717 let block2 = MemoryBlock {
1718 ptr: 2024,
1719 size: 2048,
1720 is_allocated: true,
1721 allocation_time: Some(Instant::now()),
1722 };
1723
1724 assert!(!block1.is_allocated);
1725 assert!(block2.is_allocated);
1726 assert!(block1.allocation_time.is_none());
1727 assert!(block2.allocation_time.is_some());
1728 }
1729
1730 #[test]
1731 fn test_distributed_workload() {
1732 let workload = DistributedWorkload {
1733 total_elements: 1_000_000,
1734 operation_type: "fft".to_string(),
1735 memory_requirement: 8 * 1_000_000, computation_complexity: 2.5, };
1738
1739 assert_eq!(workload.total_elements, 1_000_000);
1740 assert_eq!(workload.operation_type, "fft");
1741 assert!(workload.computation_complexity > 1.0);
1742 }
1743
1744 #[test]
1745 fn test_communication_topology() {
1746 let ring_topology = CommunicationTopology::Ring;
1747 let tree_topology = CommunicationTopology::Tree;
1748 let all_to_all_topology = CommunicationTopology::AllToAll;
1749 let custom_topology =
1750 CommunicationTopology::Custom(vec![vec![1, 2], vec![0, 3], vec![0, 3], vec![1, 2]]);
1751
1752 match ring_topology {
1754 CommunicationTopology::Ring => {}
1755 _ => panic!(),
1756 }
1757 match tree_topology {
1758 CommunicationTopology::Tree => {}
1759 _ => panic!(),
1760 }
1761 match all_to_all_topology {
1762 CommunicationTopology::AllToAll => {}
1763 _ => panic!(),
1764 }
1765 match custom_topology {
1766 CommunicationTopology::Custom(_) => {}
1767 _ => panic!(),
1768 }
1769 }
1770
1771 #[test]
1772 fn test_synchronization_barrier() {
1773 let barrier = SynchronizationBarrier {
1774 id: 1,
1775 participating_gpus: vec![0, 1, 2, 3],
1776 barrier_type: BarrierType::Global,
1777 };
1778
1779 assert_eq!(barrier.id, 1);
1780 assert_eq!(barrier.participating_gpus.len(), 4);
1781 assert!(matches!(barrier.barrier_type, BarrierType::Global));
1782 }
1783
1784 #[test]
1785 fn test_optimization_recommendation_priorities() {
1786 let low_priority = RecommendationPriority::Low;
1787 let medium_priority = RecommendationPriority::Medium;
1788 let high_priority = RecommendationPriority::High;
1789 let critical_priority = RecommendationPriority::Critical;
1790
1791 match low_priority {
1793 RecommendationPriority::Low => {}
1794 _ => panic!(),
1795 }
1796 match medium_priority {
1797 RecommendationPriority::Medium => {}
1798 _ => panic!(),
1799 }
1800 match high_priority {
1801 RecommendationPriority::High => {}
1802 _ => panic!(),
1803 }
1804 match critical_priority {
1805 RecommendationPriority::Critical => {}
1806 _ => panic!(),
1807 }
1808 }
1809
1810 #[test]
1811 fn test_performance_metric_calculations() {
1812 let metric = PerformanceMetric {
1813 execution_time: 10.0, throughput: 1000.0, memory_bandwidth: 0.8, occupancy: 0.75, };
1818
1819 assert!(metric.execution_time > 0.0);
1820 assert!(metric.throughput > 0.0);
1821 assert!(metric.memory_bandwidth <= 1.0);
1822 assert!(metric.occupancy <= 1.0);
1823 }
1824
1825 #[test]
1826 fn test_operation_status_transitions() {
1827 let mut operation = AsyncOperation {
1828 id: 0,
1829 kernel_info: GpuKernelInfo {
1830 name: "test".to_string(),
1831 device_id: 0,
1832 grid_size: (1, 1, 1),
1833 block_size: (1, 1, 1),
1834 shared_memory: 0,
1835 parameters: HashMap::new(),
1836 },
1837 stream_id: 0,
1838 start_time: Instant::now(),
1839 status: OperationStatus::Pending,
1840 };
1841
1842 assert!(matches!(operation.status, OperationStatus::Pending));
1843
1844 operation.status = OperationStatus::Running;
1845 assert!(matches!(operation.status, OperationStatus::Running));
1846
1847 operation.status = OperationStatus::Completed;
1848 assert!(matches!(operation.status, OperationStatus::Completed));
1849 }
1850}