oxigdal_gpu_advanced/multi_gpu/
sync.rs

1//! Advanced multi-GPU synchronization primitives.
2//!
3//! This module provides sophisticated synchronization mechanisms for coordinating
4//! operations across multiple GPUs including barriers, events, and cross-GPU transfers.
5
6use crate::error::{GpuAdvancedError, Result};
7use parking_lot::{Mutex, RwLock};
8use std::collections::HashMap;
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11use tokio::sync::{Notify, Semaphore};
12use wgpu::{Buffer, Device, Queue};
13
14/// Multi-GPU synchronization manager
15#[derive(Clone)]
16pub struct SyncManager {
17    devices: Arc<Vec<Arc<Device>>>,
18    queues: Arc<Vec<Arc<Queue>>>,
19    barriers: Arc<RwLock<HashMap<String, Arc<Barrier>>>>,
20    events: Arc<RwLock<HashMap<String, Arc<Event>>>>,
21    fence_pool: Arc<Mutex<FencePool>>,
22}
23
24impl SyncManager {
25    /// Create a new synchronization manager
26    pub fn new(devices: Vec<Arc<Device>>, queues: Vec<Arc<Queue>>) -> Result<Self> {
27        if devices.len() != queues.len() {
28            return Err(GpuAdvancedError::InvalidConfiguration(
29                "Device and queue count mismatch".to_string(),
30            ));
31        }
32
33        Ok(Self {
34            devices: Arc::new(devices),
35            queues: Arc::new(queues),
36            barriers: Arc::new(RwLock::new(HashMap::new())),
37            events: Arc::new(RwLock::new(HashMap::new())),
38            fence_pool: Arc::new(Mutex::new(FencePool::new())),
39        })
40    }
41
42    /// Create a barrier for N GPUs
43    pub fn create_barrier(&self, name: &str, gpu_count: usize) -> Result<Arc<Barrier>> {
44        if gpu_count == 0 || gpu_count > self.devices.len() {
45            return Err(GpuAdvancedError::ConfigError(format!(
46                "Invalid GPU count {} for barrier (available: {})",
47                gpu_count,
48                self.devices.len()
49            )));
50        }
51
52        let barrier = Arc::new(Barrier::new(gpu_count));
53        self.barriers
54            .write()
55            .insert(name.to_string(), barrier.clone());
56        Ok(barrier)
57    }
58
59    /// Get an existing barrier
60    pub fn get_barrier(&self, name: &str) -> Option<Arc<Barrier>> {
61        self.barriers.read().get(name).cloned()
62    }
63
64    /// Create an event for GPU-to-GPU synchronization
65    pub fn create_event(&self, name: &str) -> Arc<Event> {
66        let event = Arc::new(Event::new());
67        self.events.write().insert(name.to_string(), event.clone());
68        event
69    }
70
71    /// Get an existing event
72    pub fn get_event(&self, name: &str) -> Option<Arc<Event>> {
73        self.events.read().get(name).cloned()
74    }
75
76    /// Transfer data between GPUs
77    pub async fn transfer_between_gpus(
78        &self,
79        src_gpu_idx: usize,
80        dst_gpu_idx: usize,
81        src_buffer: &Buffer,
82        dst_buffer: &Buffer,
83        size: u64,
84    ) -> Result<Duration> {
85        if src_gpu_idx >= self.devices.len() || dst_gpu_idx >= self.devices.len() {
86            return Err(GpuAdvancedError::InvalidConfiguration(
87                "GPU index out of bounds".to_string(),
88            ));
89        }
90
91        let start = Instant::now();
92
93        // Create staging buffer on host
94        let staging_buffer = self.devices[src_gpu_idx].create_buffer(&wgpu::BufferDescriptor {
95            label: Some("cross_gpu_staging"),
96            size,
97            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
98            mapped_at_creation: false,
99        });
100
101        // Copy from source GPU to staging
102        let mut encoder =
103            self.devices[src_gpu_idx].create_command_encoder(&wgpu::CommandEncoderDescriptor {
104                label: Some("cross_gpu_copy_src"),
105            });
106        encoder.copy_buffer_to_buffer(src_buffer, 0, &staging_buffer, 0, size);
107        self.queues[src_gpu_idx].submit(Some(encoder.finish()));
108
109        // Wait for copy to complete
110        let slice = staging_buffer.slice(..);
111        let (tx, rx) = futures::channel::oneshot::channel();
112        slice.map_async(wgpu::MapMode::Read, move |result| {
113            let _ = tx.send(result);
114        });
115        // wgpu 28 polls automatically - no explicit poll needed
116
117        rx.await
118            .map_err(|_| GpuAdvancedError::SyncError("Transfer channel closed".to_string()))?
119            .map_err(|e| GpuAdvancedError::SyncError(format!("Map async failed: {:?}", e)))?;
120
121        // Read from staging
122        let data = slice
123            .get_mapped_range()
124            .map_err(|e| GpuAdvancedError::SyncError(format!("get_mapped_range failed: {e}")))?;
125        let vec_data: Vec<u8> = data.to_vec();
126        drop(data);
127        staging_buffer.unmap();
128
129        // Create destination staging buffer
130        let dst_staging = self.devices[dst_gpu_idx].create_buffer(&wgpu::BufferDescriptor {
131            label: Some("cross_gpu_staging_dst"),
132            size,
133            usage: wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::MAP_WRITE,
134            mapped_at_creation: true,
135        });
136
137        // Write to destination staging
138        {
139            let mut mapped = dst_staging.slice(..).get_mapped_range_mut().map_err(|e| {
140                GpuAdvancedError::SyncError(format!("get_mapped_range_mut failed: {e}"))
141            })?;
142            mapped.copy_from_slice(&vec_data);
143        }
144        dst_staging.unmap();
145
146        // Copy from staging to destination GPU
147        let mut encoder =
148            self.devices[dst_gpu_idx].create_command_encoder(&wgpu::CommandEncoderDescriptor {
149                label: Some("cross_gpu_copy_dst"),
150            });
151        encoder.copy_buffer_to_buffer(&dst_staging, 0, dst_buffer, 0, size);
152        self.queues[dst_gpu_idx].submit(Some(encoder.finish()));
153
154        // wgpu 28 polls automatically - submission completes asynchronously
155
156        Ok(start.elapsed())
157    }
158
159    /// Acquire a fence from the pool
160    pub fn acquire_fence(&self) -> Fence {
161        self.fence_pool.lock().acquire()
162    }
163
164    /// Release a fence back to the pool
165    pub fn release_fence(&self, fence: Fence) {
166        self.fence_pool.lock().release(fence);
167    }
168
169    /// Number of available GPUs
170    pub fn gpu_count(&self) -> usize {
171        self.devices.len()
172    }
173}
174
175/// Barrier synchronization primitive for multiple GPUs
176pub struct Barrier {
177    count: usize,
178    arrived: Mutex<usize>,
179    generation: Mutex<usize>,
180    notify: Notify,
181}
182
183impl Barrier {
184    /// Create a new barrier
185    pub fn new(count: usize) -> Self {
186        Self {
187            count,
188            arrived: Mutex::new(0),
189            generation: Mutex::new(0),
190            notify: Notify::new(),
191        }
192    }
193
194    /// Wait at the barrier
195    pub async fn wait(&self) -> Result<()> {
196        let current_gen = *self.generation.lock();
197
198        let arrived = {
199            let mut arrived = self.arrived.lock();
200            *arrived += 1;
201            *arrived
202        };
203
204        if arrived == self.count {
205            // Last one to arrive, reset and notify all
206            {
207                let mut arrived = self.arrived.lock();
208                *arrived = 0;
209            }
210            {
211                let mut gen_val = self.generation.lock();
212                *gen_val += 1;
213            }
214            self.notify.notify_waiters();
215            Ok(())
216        } else {
217            // Wait for notification
218            loop {
219                self.notify.notified().await;
220                let gen_val = *self.generation.lock();
221                if gen_val > current_gen {
222                    break;
223                }
224            }
225            Ok(())
226        }
227    }
228
229    /// Wait with timeout
230    pub async fn wait_timeout(&self, timeout: Duration) -> Result<bool> {
231        let wait_future = self.wait();
232        match tokio::time::timeout(timeout, wait_future).await {
233            Ok(Ok(())) => Ok(true),
234            Ok(Err(e)) => Err(e),
235            Err(_) => Ok(false), // Timeout
236        }
237    }
238
239    /// Get the barrier count
240    pub fn count(&self) -> usize {
241        self.count
242    }
243
244    /// Get the current number of waiting threads
245    pub fn waiting(&self) -> usize {
246        *self.arrived.lock()
247    }
248}
249
250/// Event for GPU-to-GPU signaling
251pub struct Event {
252    signaled: Mutex<bool>,
253    notify: Notify,
254    timestamp: Mutex<Option<Instant>>,
255}
256
257impl Event {
258    /// Create a new event
259    pub fn new() -> Self {
260        Self {
261            signaled: Mutex::new(false),
262            notify: Notify::new(),
263            timestamp: Mutex::new(None),
264        }
265    }
266
267    /// Signal the event
268    pub fn signal(&self) {
269        *self.signaled.lock() = true;
270        *self.timestamp.lock() = Some(Instant::now());
271        self.notify.notify_waiters();
272    }
273
274    /// Reset the event
275    pub fn reset(&self) {
276        *self.signaled.lock() = false;
277        *self.timestamp.lock() = None;
278    }
279
280    /// Wait for the event to be signaled
281    pub async fn wait(&self) {
282        if *self.signaled.lock() {
283            return;
284        }
285        self.notify.notified().await;
286    }
287
288    /// Wait with timeout
289    pub async fn wait_timeout(&self, timeout: Duration) -> bool {
290        if *self.signaled.lock() {
291            return true;
292        }
293        tokio::time::timeout(timeout, self.notify.notified())
294            .await
295            .is_ok()
296    }
297
298    /// Check if the event is signaled
299    pub fn is_signaled(&self) -> bool {
300        *self.signaled.lock()
301    }
302
303    /// Get the timestamp when the event was signaled
304    pub fn timestamp(&self) -> Option<Instant> {
305        *self.timestamp.lock()
306    }
307}
308
309impl Default for Event {
310    fn default() -> Self {
311        Self::new()
312    }
313}
314
315/// Fence for command buffer synchronization
316#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
317pub struct Fence {
318    id: u64,
319}
320
321impl Fence {
322    fn new(id: u64) -> Self {
323        Self { id }
324    }
325
326    /// Get the fence ID
327    pub fn id(&self) -> u64 {
328        self.id
329    }
330}
331
332/// Pool of reusable fences
333struct FencePool {
334    next_id: u64,
335    available: Vec<Fence>,
336    max_pool_size: usize,
337}
338
339impl FencePool {
340    fn new() -> Self {
341        Self {
342            next_id: 0,
343            available: Vec::new(),
344            max_pool_size: 256,
345        }
346    }
347
348    fn acquire(&mut self) -> Fence {
349        if let Some(fence) = self.available.pop() {
350            fence
351        } else {
352            let fence = Fence::new(self.next_id);
353            self.next_id += 1;
354            fence
355        }
356    }
357
358    fn release(&mut self, fence: Fence) {
359        if self.available.len() < self.max_pool_size {
360            self.available.push(fence);
361        }
362    }
363}
364
365/// Semaphore for controlling concurrent GPU access
366pub struct GpuSemaphore {
367    inner: Arc<Semaphore>,
368}
369
370impl GpuSemaphore {
371    /// Create a new semaphore with the given permit count
372    pub fn new(permits: usize) -> Self {
373        Self {
374            inner: Arc::new(Semaphore::new(permits)),
375        }
376    }
377
378    /// Acquire a permit
379    pub async fn acquire(&self) -> Result<SemaphoreGuard<'_>> {
380        let permit =
381            self.inner.acquire().await.map_err(|e| {
382                GpuAdvancedError::SyncError(format!("Semaphore acquire failed: {}", e))
383            })?;
384        Ok(SemaphoreGuard { _permit: permit })
385    }
386
387    /// Try to acquire a permit without waiting
388    pub fn try_acquire(&self) -> Option<SemaphoreGuard<'_>> {
389        self.inner
390            .try_acquire()
391            .ok()
392            .map(|permit| SemaphoreGuard { _permit: permit })
393    }
394
395    /// Get available permits
396    pub fn available_permits(&self) -> usize {
397        self.inner.available_permits()
398    }
399}
400
401impl Clone for GpuSemaphore {
402    fn clone(&self) -> Self {
403        Self {
404            inner: Arc::clone(&self.inner),
405        }
406    }
407}
408
409/// RAII guard for semaphore permit
410pub struct SemaphoreGuard<'a> {
411    _permit: tokio::sync::SemaphorePermit<'a>,
412}
413
414/// Synchronization statistics
415#[derive(Debug, Clone, Default)]
416pub struct SyncStats {
417    /// Number of barrier waits
418    pub barrier_waits: u64,
419    /// Number of event signals
420    pub event_signals: u64,
421    /// Number of cross-GPU transfers
422    pub cross_gpu_transfers: u64,
423    /// Total transfer time
424    pub total_transfer_time: Duration,
425    /// Total bytes transferred
426    pub total_bytes_transferred: u64,
427}
428
429impl SyncStats {
430    /// Calculate average transfer bandwidth in GB/s
431    pub fn average_bandwidth_gbs(&self) -> Option<f64> {
432        if self.total_transfer_time > Duration::ZERO && self.total_bytes_transferred > 0 {
433            let bytes_per_sec =
434                self.total_bytes_transferred as f64 / self.total_transfer_time.as_secs_f64();
435            Some(bytes_per_sec / 1_000_000_000.0)
436        } else {
437            None
438        }
439    }
440}
441
442#[cfg(test)]
443mod tests {
444    use super::*;
445
446    #[tokio::test]
447    async fn test_barrier() {
448        let barrier = Arc::new(Barrier::new(3));
449        let mut handles = Vec::new();
450
451        for i in 0..3 {
452            let b = barrier.clone();
453            let handle = tokio::spawn(async move {
454                println!("Task {} waiting at barrier", i);
455                b.wait().await.ok();
456                println!("Task {} passed barrier", i);
457            });
458            handles.push(handle);
459        }
460
461        for handle in handles {
462            handle.await.ok();
463        }
464
465        assert_eq!(barrier.waiting(), 0);
466    }
467
468    #[tokio::test]
469    async fn test_event() {
470        let event = Arc::new(Event::new());
471        assert!(!event.is_signaled());
472
473        let e = event.clone();
474        let handle = tokio::spawn(async move {
475            e.wait().await;
476        });
477
478        tokio::time::sleep(Duration::from_millis(10)).await;
479        event.signal();
480        assert!(event.is_signaled());
481
482        handle.await.ok();
483    }
484
485    #[tokio::test]
486    async fn test_semaphore() {
487        let sem = GpuSemaphore::new(2);
488        assert_eq!(sem.available_permits(), 2);
489
490        let _guard1 = sem.acquire().await.ok();
491        assert_eq!(sem.available_permits(), 1);
492
493        let _guard2 = sem.acquire().await.ok();
494        assert_eq!(sem.available_permits(), 0);
495
496        drop(_guard1);
497        assert_eq!(sem.available_permits(), 1);
498    }
499
500    #[test]
501    fn test_fence_pool() {
502        let mut pool = FencePool::new();
503        let f1 = pool.acquire();
504        let f2 = pool.acquire();
505
506        assert_ne!(f1.id(), f2.id());
507
508        pool.release(f1);
509        let f3 = pool.acquire();
510        assert_eq!(f1, f3);
511    }
512}
oxigdal_gpu_advanced/multi_gpu/sync.rs

oxigdal_gpu_advanced/multi_gpu/
sync.rs
