1use crate::kernels;
6use crate::oneshot;
7use flashraw_core::{Error, Result};
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::{Arc, Mutex, OnceLock};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum AdapterClass {
15 Discrete,
16 Integrated,
18 Virtual,
19 Software,
21 Unknown,
23}
24
25impl AdapterClass {
26 pub fn name(self) -> &'static str {
27 match self {
28 AdapterClass::Discrete => "discrete",
29 AdapterClass::Integrated => "integrated",
30 AdapterClass::Virtual => "virtual",
31 AdapterClass::Software => "software",
32 AdapterClass::Unknown => "unknown",
33 }
34 }
35}
36
37#[derive(Debug, Clone, PartialEq)]
39pub struct GpuAdapterInfo {
40 pub name: String,
41 pub backend: String,
42 pub driver: String,
43 pub class: AdapterClass,
44 pub timestamps: bool,
46 pub max_buffer_size: u64,
47 pub max_storage_binding: u64,
48}
49
50impl std::fmt::Display for GpuAdapterInfo {
51 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52 write!(f, "{} ({}, {})", self.name, self.backend, self.class.name())
53 }
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
58pub(crate) enum Kernel {
59 Unpack,
60 Normalize,
61 Develop,
62 DevelopHalf,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
67pub(crate) struct PipelineKey {
68 pub kernel: Kernel,
69 pub input_mode: u32,
70 pub order: u32,
71 pub use_lut: bool,
72 pub layout: u32,
74 pub out_mode: u32,
75}
76
77pub struct GpuContext {
83 pub(crate) adapter: Option<wgpu::Adapter>,
84 pub(crate) device: wgpu::Device,
85 pub(crate) queue: wgpu::Queue,
86 pub(crate) info: GpuAdapterInfo,
87 pub(crate) layout: wgpu::BindGroupLayout,
88 pipeline_layout: wgpu::PipelineLayout,
89 pub(crate) tex_layout: wgpu::BindGroupLayout,
90 tex_pipeline: OnceLock<wgpu::ComputePipeline>,
91 modules: Mutex<HashMap<Kernel, wgpu::ShaderModule>>,
92 pipelines: Mutex<HashMap<PipelineKey, wgpu::ComputePipeline>>,
93 pub(crate) srgb: Vec<u32>,
94 pub(crate) last_error: Arc<Mutex<Option<String>>>,
95 pub(crate) calibration: OnceLock<crate::policy::GpuCalibration>,
96 pub(crate) uma: bool,
98 pool: Mutex<Vec<wgpu::Buffer>>,
100 binding_cap: AtomicU64,
102 buffer_cap: AtomicU64,
103}
104
105const POOL_MAX_BYTES: u64 = 1 << 30;
107
108impl std::fmt::Debug for GpuContext {
109 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110 f.debug_struct("GpuContext").field("adapter", &self.info).finish()
111 }
112}
113
114fn classify(t: wgpu::DeviceType) -> AdapterClass {
115 match t {
116 wgpu::DeviceType::DiscreteGpu => AdapterClass::Discrete,
117 wgpu::DeviceType::IntegratedGpu => AdapterClass::Integrated,
118 wgpu::DeviceType::VirtualGpu => AdapterClass::Virtual,
119 wgpu::DeviceType::Cpu => AdapterClass::Software,
120 wgpu::DeviceType::Other => AdapterClass::Unknown,
121 }
122}
123
124fn is_software_name(name: &str) -> bool {
125 let n = name.to_ascii_lowercase();
126 ["llvmpipe", "lavapipe", "softpipe", "swiftshader", "microsoft basic render", "warp"].iter().any(|s| n.contains(s))
127}
128
129fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
130 wgpu::BindGroupLayoutEntry {
131 binding,
132 visibility: wgpu::ShaderStages::COMPUTE,
133 ty: wgpu::BindingType::Buffer {
134 ty: wgpu::BufferBindingType::Storage { read_only },
135 has_dynamic_offset: false,
136 min_binding_size: None,
137 },
138 count: None,
139 }
140}
141
142fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
143 wgpu::BindGroupLayoutEntry {
144 binding,
145 visibility: wgpu::ShaderStages::COMPUTE,
146 ty: wgpu::BindingType::Buffer {
147 ty: wgpu::BufferBindingType::Uniform,
148 has_dynamic_offset: false,
149 min_binding_size: None,
150 },
151 count: None,
152 }
153}
154
155impl GpuContext {
156 pub async fn new() -> Result<Self> {
158 if disabled_by_env() {
159 return Err(Error::GpuUnavailable("disabled by `FLASHRAW_GPU`".into()));
160 }
161 let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle_from_env());
162 let adapter = instance
163 .request_adapter(&wgpu::RequestAdapterOptions {
164 power_preference: wgpu::PowerPreference::HighPerformance,
165 force_fallback_adapter: false,
166 compatible_surface: None,
167 apply_limit_buckets: false,
168 })
169 .await
170 .map_err(|e| Error::GpuUnavailable(format!("no GPU adapter: {e}")))?;
171 let mut features = wgpu::Features::empty();
172 if adapter.features().contains(wgpu::Features::TIMESTAMP_QUERY) {
173 features |= wgpu::Features::TIMESTAMP_QUERY;
174 }
175 let integrated = adapter.get_info().device_type == wgpu::DeviceType::IntegratedGpu;
178 if integrated && adapter.features().contains(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS) {
179 features |= wgpu::Features::MAPPABLE_PRIMARY_BUFFERS;
180 }
181 let limits = adapter.limits();
183 let (device, queue) = adapter
184 .request_device(&wgpu::DeviceDescriptor {
185 label: Some("flashraw"),
186 required_features: features,
187 required_limits: limits,
188 experimental_features: wgpu::ExperimentalFeatures::default(),
189 memory_hints: wgpu::MemoryHints::Performance,
190 trace: wgpu::Trace::Off,
191 })
192 .await
193 .map_err(|e| Error::GpuUnavailable(format!("device request failed: {e}")))?;
194 let info = adapter.get_info();
195 Ok(Self::build(Some(adapter), device, queue, &info))
196 }
197
198 pub fn from_device(device: wgpu::Device, queue: wgpu::Queue, info: &wgpu::AdapterInfo) -> Self {
202 Self::build(None, device, queue, info)
203 }
204
205 fn build(adapter: Option<wgpu::Adapter>, device: wgpu::Device, queue: wgpu::Queue, ai: &wgpu::AdapterInfo) -> Self {
206 let mut class = classify(ai.device_type);
207 if is_software_name(&ai.name) {
208 class = AdapterClass::Software;
209 }
210 let limits = device.limits();
211 let uma = device.features().contains(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS);
212 let info = GpuAdapterInfo {
213 name: if ai.name.is_empty() { "WebGPU adapter".into() } else { ai.name.clone() },
214 backend: format!("{:?}", ai.backend),
215 driver: [ai.driver.as_str(), ai.driver_info.as_str()]
216 .iter()
217 .filter(|s| !s.is_empty())
218 .cloned()
219 .collect::<Vec<_>>()
220 .join(" "),
221 class,
222 timestamps: device.features().contains(wgpu::Features::TIMESTAMP_QUERY),
223 max_buffer_size: limits.max_buffer_size,
224 max_storage_binding: limits.max_storage_buffer_binding_size,
225 };
226 let last_error: Arc<Mutex<Option<String>>> = Arc::default();
227 {
228 let slot = last_error.clone();
231 device.on_uncaptured_error(Arc::new(move |e: wgpu::Error| {
232 if let Ok(mut s) = slot.lock() {
233 *s = Some(e.to_string());
234 }
235 }));
236 }
237 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
238 label: Some("flashraw.kernels"),
239 entries: &[
240 uniform_entry(0),
241 storage_entry(1, true),
242 storage_entry(2, true),
243 storage_entry(3, true),
244 storage_entry(4, false),
245 ],
246 });
247 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
248 label: Some("flashraw.kernels"),
249 bind_group_layouts: &[Some(&layout)],
250 immediate_size: 0,
251 });
252 let tex_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
253 label: Some("flashraw.to_texture"),
254 entries: &[
255 uniform_entry(0),
256 storage_entry(1, true),
257 wgpu::BindGroupLayoutEntry {
258 binding: 2,
259 visibility: wgpu::ShaderStages::COMPUTE,
260 ty: wgpu::BindingType::StorageTexture {
261 access: wgpu::StorageTextureAccess::WriteOnly,
262 format: wgpu::TextureFormat::Rgba8Unorm,
263 view_dimension: wgpu::TextureViewDimension::D2,
264 },
265 count: None,
266 },
267 ],
268 });
269 GpuContext {
270 adapter,
271 device,
272 queue,
273 info,
274 layout,
275 pipeline_layout,
276 tex_layout,
277 tex_pipeline: OnceLock::new(),
278 modules: Mutex::default(),
279 pipelines: Mutex::default(),
280 srgb: kernels::srgb_table(),
281 last_error,
282 calibration: OnceLock::new(),
283 uma,
284 pool: Mutex::default(),
285 binding_cap: AtomicU64::new(0),
286 buffer_cap: AtomicU64::new(0),
287 }
288 }
289
290 pub(crate) fn acquire(&self, label: &str, size: u64, usage: wgpu::BufferUsages) -> wgpu::Buffer {
293 let size = size.max(4);
294 {
295 let mut pool = self.pool.lock().unwrap_or_else(|e| e.into_inner());
296 let best = pool
297 .iter()
298 .enumerate()
299 .filter(|(_, b)| b.usage() == usage && b.size() >= size && b.size() <= size + size / 4 + (16 << 20))
300 .min_by_key(|(_, b)| b.size())
301 .map(|(i, _)| i);
302 if let Some(i) = best {
303 return pool.swap_remove(i);
304 }
305 }
306 let granule: u64 = if size > (64 << 20) { 4 << 20 } else { 256 };
307 let rounded = (size.div_ceil(granule) * granule).min(self.info.max_buffer_size & !3).max(size);
309 self.device.create_buffer(&wgpu::BufferDescriptor {
310 label: Some(label),
311 size: rounded,
312 usage,
313 mapped_at_creation: false,
314 })
315 }
316
317 pub(crate) fn release(&self, b: wgpu::Buffer) {
319 if b.size() < (1 << 20) {
320 return; }
322 let mut pool = self.pool.lock().unwrap_or_else(|e| e.into_inner());
323 pool.push(b);
324 let mut total: u64 = pool.iter().map(|b| b.size()).sum();
325 while total > POOL_MAX_BYTES && !pool.is_empty() {
326 let old = pool.remove(0);
327 total -= old.size();
328 old.destroy();
329 }
330 }
331
332 pub fn trim(&self) {
334 let mut pool = self.pool.lock().unwrap_or_else(|e| e.into_inner());
335 for b in pool.drain(..) {
336 b.destroy();
337 }
338 }
339
340 pub fn unified_memory(&self) -> bool {
342 self.uma
343 }
344
345 pub fn info(&self) -> &GpuAdapterInfo {
346 &self.info
347 }
348 pub fn device(&self) -> &wgpu::Device {
349 &self.device
350 }
351 pub fn queue(&self) -> &wgpu::Queue {
352 &self.queue
353 }
354 pub fn adapter(&self) -> Option<&wgpu::Adapter> {
356 self.adapter.as_ref()
357 }
358
359 fn module(&self, k: Kernel) -> wgpu::ShaderModule {
360 let mut m = self.modules.lock().unwrap_or_else(|e| e.into_inner());
361 m.entry(k)
362 .or_insert_with(|| {
363 let (label, src) = match k {
364 Kernel::Unpack => ("flashraw.unpack", kernels::UNPACK_WGSL),
365 Kernel::Normalize => ("flashraw.normalize", kernels::NORMALIZE_WGSL),
366 Kernel::Develop | Kernel::DevelopHalf => ("flashraw.develop", kernels::DEVELOP_WGSL),
367 };
368 self.device.create_shader_module(wgpu::ShaderModuleDescriptor {
369 label: Some(label),
370 source: wgpu::ShaderSource::Wgsl(src.into()),
371 })
372 })
373 .clone()
374 }
375
376 pub(crate) fn pipeline(&self, key: PipelineKey) -> wgpu::ComputePipeline {
378 if let Some(p) = self.pipelines.lock().unwrap_or_else(|e| e.into_inner()).get(&key) {
379 return p.clone();
380 }
381 let module = self.module(key.kernel);
382 let entry = match key.kernel {
383 Kernel::Unpack => "unpack",
384 Kernel::Normalize => "normalize",
385 Kernel::Develop => "develop",
386 Kernel::DevelopHalf => "develop_half",
387 };
388 let mut constants: Vec<(&str, f64)> = vec![
389 ("INPUT_MODE", key.input_mode as f64),
390 ("ORDER", key.order as f64),
391 ("USE_LUT", if key.use_lut { 1.0 } else { 0.0 }),
392 ];
393 if matches!(key.kernel, Kernel::Develop | Kernel::DevelopHalf) {
394 constants.push(("LAYOUT", key.layout as f64));
395 constants.push(("OUT_MODE", key.out_mode as f64));
396 }
397 let p = self.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
398 label: Some(entry),
399 layout: Some(&self.pipeline_layout),
400 module: &module,
401 entry_point: Some(entry),
402 compilation_options: wgpu::PipelineCompilationOptions {
403 constants: &constants,
404 zero_initialize_workgroup_memory: false,
406 },
407 cache: None,
408 });
409 self.pipelines.lock().unwrap_or_else(|e| e.into_inner()).insert(key, p.clone());
410 p
411 }
412
413 pub(crate) fn tex_pipeline(&self) -> &wgpu::ComputePipeline {
414 self.tex_pipeline.get_or_init(|| {
415 let module = self.device.create_shader_module(wgpu::ShaderModuleDescriptor {
416 label: Some("flashraw.to_texture"),
417 source: wgpu::ShaderSource::Wgsl(kernels::TO_TEXTURE_WGSL.into()),
418 });
419 let layout = self.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
420 label: Some("flashraw.to_texture"),
421 bind_group_layouts: &[Some(&self.tex_layout)],
422 immediate_size: 0,
423 });
424 self.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
425 label: Some("to_texture"),
426 layout: Some(&layout),
427 module: &module,
428 entry_point: Some("to_texture"),
429 compilation_options: Default::default(),
430 cache: None,
431 })
432 })
433 }
434
435 pub async fn wait_idle(&self) -> Result<()> {
438 let (tx, rx) = oneshot::channel();
439 self.queue.on_submitted_work_done(move || tx.send(()));
440 #[cfg(not(target_arch = "wasm32"))]
441 {
442 self.device
443 .poll(wgpu::PollType::wait_indefinitely())
444 .map_err(|e| Error::GpuUnavailable(format!("device poll failed: {e}")))?;
445 if rx.try_take().is_some() {
446 return Ok(());
447 }
448 }
449 rx.await.ok_or_else(|| Error::GpuUnavailable("GPU work did not complete (device lost?)".into()))
450 }
451
452 pub(crate) async fn read_mapped<T>(
454 &self,
455 buffer: &wgpu::Buffer,
456 size: u64,
457 f: impl FnOnce(&[u8]) -> T,
458 ) -> Result<T> {
459 let (tx, rx) = oneshot::channel();
460 buffer.map_async(wgpu::MapMode::Read, ..size, move |r| tx.send(r));
461 #[cfg(not(target_arch = "wasm32"))]
462 self.device
463 .poll(wgpu::PollType::wait_indefinitely())
464 .map_err(|e| Error::GpuUnavailable(format!("device poll failed: {e}")))?;
465 let r = match rx.try_take() {
466 Some(r) => Some(r),
467 None => rx.await,
468 };
469 match r {
470 Some(Ok(())) => {}
471 Some(Err(e)) => return Err(Error::GpuUnavailable(format!("buffer mapping failed: {e}"))),
472 None => return Err(Error::GpuUnavailable("buffer mapping was cancelled".into())),
473 }
474 let out = {
475 let view = buffer
476 .get_mapped_range(..size)
477 .map_err(|e| Error::GpuUnavailable(format!("mapped range unavailable: {e}")))?;
478 f(&view)
479 };
480 buffer.unmap();
481 Ok(out)
482 }
483
484 pub(crate) async fn write_mapped(
491 &self,
492 buffer: &wgpu::Buffer,
493 data: &[u8],
494 after: Option<wgpu::SubmissionIndex>,
495 ) -> Result<f64> {
496 let size = (data.len() as u64).div_ceil(4).max(1) * 4;
497 let (tx, rx) = oneshot::channel();
498 buffer.map_async(wgpu::MapMode::Write, ..size, move |r| tx.send(r));
499 #[cfg(not(target_arch = "wasm32"))]
500 {
501 let poll = match after {
502 Some(i) => wgpu::PollType::Wait { submission_index: Some(i), timeout: None },
503 None => wgpu::PollType::Poll,
504 };
505 self.device.poll(poll).map_err(|e| Error::GpuUnavailable(format!("device poll failed: {e}")))?;
506 if !rx.is_ready() {
507 self.device
508 .poll(wgpu::PollType::wait_indefinitely())
509 .map_err(|e| Error::GpuUnavailable(format!("device poll failed: {e}")))?;
510 }
511 }
512 #[cfg(target_arch = "wasm32")]
513 let _ = after;
514 let r = match rx.try_take() {
515 Some(r) => Some(r),
516 None => rx.await,
517 };
518 match r {
519 Some(Ok(())) => {}
520 Some(Err(e)) => return Err(Error::GpuUnavailable(format!("buffer mapping failed: {e}"))),
521 None => return Err(Error::GpuUnavailable("buffer mapping was cancelled".into())),
522 }
523 let t0 = flashraw_core::profile::now_ms();
524 {
525 let mut view = buffer
526 .get_mapped_range_mut(..size)
527 .map_err(|e| Error::GpuUnavailable(format!("mapped range unavailable: {e}")))?;
528 let bulk = data.len();
529 view.slice(..bulk).copy_from_slice(data);
530 if (bulk as u64) < size {
531 view.slice(bulk..).fill(0);
532 }
533 }
534 buffer.unmap();
535 Ok(flashraw_core::profile::now_ms() - t0)
536 }
537
538 pub(crate) fn take_error(&self) -> Option<String> {
540 self.last_error.lock().ok()?.take()
541 }
542
543 pub fn override_limits(&self, max_binding: Option<u64>, max_buffer: Option<u64>) {
549 let v = |o: Option<u64>| o.map(|b| (b & !255).max(4096)).unwrap_or(0);
550 self.binding_cap.store(v(max_binding), Ordering::Relaxed);
551 self.buffer_cap.store(v(max_buffer), Ordering::Relaxed);
552 }
553
554 pub fn max_binding_bytes(&self) -> u64 {
556 let mut b = self.info.max_storage_binding.min(self.info.max_buffer_size).min(u32::MAX as u64 & !255);
557 let cap = self.binding_cap.load(Ordering::Relaxed);
558 if cap > 0 {
559 b = b.min(cap);
560 }
561 b.min(self.max_buffer_bytes())
562 }
563
564 pub fn max_buffer_bytes(&self) -> u64 {
566 let cap = self.buffer_cap.load(Ordering::Relaxed);
567 if cap > 0 { self.info.max_buffer_size.min(cap) } else { self.info.max_buffer_size }
568 }
569
570 pub fn estimate_bands(&self, input_bytes: u64, output_bytes: u64) -> u32 {
573 let b = self.max_binding_bytes().max(1);
574 let out = if output_bytes <= b { 1 } else { output_bytes.div_ceil(b) };
575 input_bytes.div_ceil(b).max(out).max(1).min(u32::MAX as u64) as u32
576 }
577}
578
579fn disabled_by_env() -> bool {
581 #[cfg(not(target_arch = "wasm32"))]
582 {
583 matches!(
584 std::env::var("FLASHRAW_GPU").map(|v| v.to_ascii_lowercase()).as_deref(),
585 Ok("0" | "off" | "false" | "none" | "cpu")
586 )
587 }
588 #[cfg(target_arch = "wasm32")]
589 {
590 false
591 }
592}
593
594#[derive(Default)]
597struct Slot {
598 ctx: Option<Arc<GpuContext>>,
599 error: Option<String>,
600 tried: bool,
601}
602
603#[cfg(not(target_arch = "wasm32"))]
604fn with_slot<T>(f: impl FnOnce(&mut Slot) -> T) -> T {
605 static SLOT: OnceLock<Mutex<Slot>> = OnceLock::new();
606 let mut g = SLOT.get_or_init(Default::default).lock().unwrap_or_else(|e| e.into_inner());
607 f(&mut g)
608}
609
610#[cfg(target_arch = "wasm32")]
611fn with_slot<T>(f: impl FnOnce(&mut Slot) -> T) -> T {
612 thread_local! { static SLOT: std::cell::RefCell<Slot> = Default::default(); }
614 SLOT.with(|s| f(&mut s.borrow_mut()))
615}
616
617pub async fn init() -> Result<Arc<GpuContext>> {
620 if let Some(r) = with_slot(|s| s.tried.then(|| s.ctx.clone().ok_or_else(|| s.error.clone().unwrap_or_default()))) {
621 return r.map_err(Error::GpuUnavailable);
622 }
623 let r = GpuContext::new().await.map(Arc::new);
624 with_slot(|s| {
625 if let Some(c) = &s.ctx {
626 return Ok(c.clone()); }
628 s.tried = true;
629 match r {
630 Ok(c) => {
631 s.ctx = Some(c.clone());
632 Ok(c)
633 }
634 Err(e) => {
635 s.error = Some(e.to_string());
636 Err(e)
637 }
638 }
639 })
640}
641
642pub fn context() -> Result<Arc<GpuContext>> {
645 #[cfg(not(target_arch = "wasm32"))]
646 {
647 if let Some(r) =
648 with_slot(|s| s.tried.then(|| s.ctx.clone().ok_or_else(|| s.error.clone().unwrap_or_default())))
649 {
650 return r.map_err(Error::GpuUnavailable);
651 }
652 pollster::block_on(init())
653 }
654 #[cfg(target_arch = "wasm32")]
655 {
656 with_slot(|s| s.ctx.clone())
657 .ok_or_else(|| Error::GpuUnavailable("call `flashraw_gpu::init().await` before using the GPU".into()))
658 }
659}
660
661pub fn try_context() -> Option<Arc<GpuContext>> {
663 with_slot(|s| s.ctx.clone())
664}
665
666pub fn set_context(ctx: Arc<GpuContext>) {
669 with_slot(|s| {
670 s.ctx = Some(ctx);
671 s.tried = true;
672 s.error = None;
673 })
674}
675
676pub fn available() -> bool {
678 context().is_ok()
679}
680
681pub fn adapter_info() -> Option<GpuAdapterInfo> {
683 context().ok().map(|c| c.info.clone())
684}