subspace_farmer/plotter/gpu/
wgpu.rs

1//! wgpu GPU records encoder
2
3use crate::plotter::gpu::GpuRecordsEncoder;
4use async_lock::Mutex as AsyncMutex;
5use parking_lot::Mutex;
6use rayon::{ThreadPool, ThreadPoolBuildError, ThreadPoolBuilder, current_thread_index};
7use std::fmt;
8use std::process::exit;
9use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, Ordering};
11use subspace_core_primitives::pieces::{PieceOffset, Record};
12use subspace_core_primitives::sectors::SectorId;
13use subspace_farmer_components::plotting::RecordsEncoder;
14use subspace_farmer_components::sector::SectorContentsMap;
15use subspace_proof_of_space_wgpu::WgpuDevice;
16
17/// wgpu implementation of [`GpuRecordsEncoder`]
18pub struct WgpuRecordsEncoder {
19    devices: Vec<Mutex<WgpuDevice>>,
20    thread_pool: ThreadPool,
21    global_mutex: Arc<AsyncMutex<()>>,
22}
23
24impl fmt::Debug for WgpuRecordsEncoder {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        f.debug_struct("WgpuRecordsEncoder").finish_non_exhaustive()
27    }
28}
29
30impl GpuRecordsEncoder for WgpuRecordsEncoder {
31    const TYPE: &'static str = "wgpu";
32}
33
34impl RecordsEncoder for WgpuRecordsEncoder {
35    fn encode_records(
36        &mut self,
37        sector_id: &SectorId,
38        records: &mut [Record],
39        abort_early: &AtomicBool,
40    ) -> anyhow::Result<SectorContentsMap> {
41        let pieces_in_sector = records
42            .len()
43            .try_into()
44            .map_err(|error| anyhow::anyhow!("Failed to convert pieces in sector: {error}"))?;
45        let mut sector_contents_map = SectorContentsMap::new(pieces_in_sector);
46
47        {
48            let iter = Mutex::new(
49                (PieceOffset::ZERO..)
50                    .zip(records.iter_mut())
51                    .zip(sector_contents_map.iter_record_bitfields_mut()),
52            );
53            let plotting_error = Mutex::new(None::<String>);
54
55            self.thread_pool.scope(|scope| {
56                scope.spawn_broadcast(|_scope, _ctx| {
57                    // One device (GPU queue) per pool thread, so this lock is always uncontended
58                    let thread_index = current_thread_index().unwrap_or_default();
59                    let Some(device) = self.devices.get(thread_index) else {
60                        return;
61                    };
62                    let mut device = device
63                        .try_lock()
64                        .expect("1:1 mapping between threads and devices; qed");
65
66                    loop {
67                        // Take mutex briefly to make sure encoding is allowed right now
68                        self.global_mutex.lock_blocking();
69
70                        // This instead of `while` above because otherwise mutex will be held for the
71                        // duration of the loop and will limit concurrency to 1 record
72                        let Some(((piece_offset, record), mut encoded_chunks_used)) =
73                            iter.lock().next()
74                        else {
75                            return;
76                        };
77                        let pos_seed = sector_id.derive_evaluation_seed(piece_offset);
78
79                        if let Err(error) = device.generate_and_encode_pospace(
80                            &pos_seed,
81                            record,
82                            encoded_chunks_used.iter_mut(),
83                        ) {
84                            plotting_error.lock().replace(error);
85                            return;
86                        }
87
88                        if abort_early.load(Ordering::Relaxed) {
89                            return;
90                        }
91                    }
92                });
93            });
94
95            let plotting_error = plotting_error.lock().take();
96            if let Some(error) = plotting_error {
97                return Err(anyhow::Error::msg(error));
98            }
99        }
100
101        Ok(sector_contents_map)
102    }
103}
104
105impl WgpuRecordsEncoder {
106    /// Create new instance.
107    ///
108    /// One thread is spawned per device (GPU queue), so records encode concurrently across queues.
109    pub fn new(
110        id: u32,
111        devices: Vec<WgpuDevice>,
112        global_mutex: Arc<AsyncMutex<()>>,
113    ) -> Result<Self, ThreadPoolBuildError> {
114        let thread_name = move |thread_index| format!("wgpu-{id:02}.{thread_index:02}");
115        // TODO: remove this panic handler when rayon logs panic_info
116        // https://github.com/rayon-rs/rayon/issues/1208
117        let panic_handler = move |panic_info| {
118            if let Some(index) = current_thread_index() {
119                eprintln!("panic on thread {}: {:?}", thread_name(index), panic_info);
120            } else {
121                // We want to guarantee exit, rather than panicking in a panic handler.
122                eprintln!("rayon panic handler called on non-rayon thread: {panic_info:?}");
123            }
124            exit(1);
125        };
126
127        let thread_pool = ThreadPoolBuilder::new()
128            .thread_name(thread_name)
129            .panic_handler(panic_handler)
130            .num_threads(devices.len())
131            .build()?;
132
133        Ok(Self {
134            devices: devices.into_iter().map(Mutex::new).collect(),
135            thread_pool,
136            global_mutex,
137        })
138    }
139}
subspace_farmer/plotter/gpu/wgpu.rs

subspace_farmer/plotter/gpu/
wgpu.rs
