zksync-gpu-prover 0.155.10

ZKsync GPU prover utilities
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use super::*;

#[derive(Debug)]
pub struct Event {
    pub(crate) sub_events: Mutex<Vec<Vec<(usize, Arc<bc_event>)>>>,
}

impl Clone for Event {
    fn clone(&self) -> Self {
        let inner = self.sub_events.lock().unwrap();
        Self {
            sub_events: Mutex::new(inner.clone()),
        }
    }
}

impl Event {
    pub fn new() -> Self {
        Self {
            sub_events: Mutex::new(vec![]),
        }
    }

    pub fn record(&self, stream: &Stream) -> GpuResult<()> {
        let device_id = stream.device_id();
        let stream_id = stream.inner.handle as usize;
        set_device(device_id)?;

        let mut sub_events = self.sub_events.lock().unwrap();

        if sub_events.len() < device_id + 1 {
            sub_events.resize(device_id + 1, vec![]);
        }

        let mut event = bc_event {
            handle: std::ptr::null_mut() as *mut c_void,
        };
        unsafe {
            let result = bc_event_create(addr_of_mut!(event), true, true);
            if result != 0 {
                return Err(GpuError::EventCreateErr(result));
            }
        }
        unsafe {
            let result = bc_event_record(event, stream.inner);
            if result != 0 {
                return Err(GpuError::EventRecordErr(result));
            }
        }

        let event = Arc::new(event);
        let inner = sub_events[device_id]
            .iter_mut()
            .filter(|(id, _)| *id == stream_id)
            .next();
        if let Some(inner) = inner {
            if Arc::strong_count(&inner.1) == 1 {
                unsafe {
                    let result = bc_event_destroy(*inner.1.as_ref());
                    if result != 0 {
                        panic!("EventDestroyErr({}) while droping Event", result);
                    }
                }
            }

            inner.1 = event;
        } else {
            sub_events[device_id].push((stream_id, event));
        }

        Ok(())
    }

    pub fn sync(&self) -> GpuResult<()> {
        let mut sub_events = self.sub_events.lock().unwrap();
        for events in sub_events.iter() {
            for (_, event) in events.iter() {
                unsafe {
                    let result = bc_event_synchronize(*event.as_ref());
                    if result != 0 {
                        return Err(GpuError::EventSyncErr(result));
                    }
                }
            }
        }

        Ok(())
    }
}

impl Drop for Event {
    fn drop(&mut self) {
        let mut sub_events = self.sub_events.lock().unwrap();
        for events in sub_events.iter() {
            for (_, event) in events.iter() {
                if Arc::strong_count(event) == 1 {
                    unsafe {
                        let result = bc_event_destroy(*event.as_ref());
                        if result != 0 {
                            panic!("EventDestroyErr({}) while droping Event", result);
                        }
                    }
                }
            }
        }
    }
}