1mod activations;
13mod backward;
14mod eigen;
15pub(crate) mod linalg;
16mod reductions;
17
18#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
19use super::runtime;
20
21#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
34static DEVICE_INIT_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
35
36#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
58pub(crate) const fn gpu_backends() -> wgpu::Backends {
59 wgpu::Backends::PRIMARY
61}
62
63#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
75pub(crate) fn shared_instance() -> wgpu::Instance {
76 use std::sync::OnceLock;
77 static INSTANCE: OnceLock<wgpu::Instance> = OnceLock::new();
78 INSTANCE
79 .get_or_init(|| {
80 wgpu::Instance::new(&wgpu::InstanceDescriptor {
94 backends: gpu_backends(),
95 ..Default::default()
96 })
97 })
98 .clone()
99}
100
101#[derive(Clone)]
103pub struct GpuDevice {
104 pub device: wgpu::Device,
105 pub queue: wgpu::Queue,
106}
107
108impl GpuDevice {
109 #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
111 pub fn new() -> Result<Self, String> {
112 let _guard = DEVICE_INIT_LOCK.lock().unwrap_or_else(|e| e.into_inner());
115 runtime::block_on(async { Self::new_async().await })
116 }
117
118 pub async fn new_async() -> Result<Self, String> {
120 let instance = shared_instance();
122
123 let adapter = instance
125 .request_adapter(&wgpu::RequestAdapterOptions {
126 power_preference: wgpu::PowerPreference::HighPerformance,
127 compatible_surface: None,
128 force_fallback_adapter: false,
129 })
130 .await
131 .map_err(|e| format!("Failed to find GPU adapter: {}", e))?;
132
133 let mut limits = wgpu::Limits::default();
137 limits.max_buffer_size = adapter.limits().max_buffer_size;
138 limits.max_storage_buffer_binding_size = adapter.limits().max_storage_buffer_binding_size;
139
140 let (device, queue) = adapter
141 .request_device(&wgpu::DeviceDescriptor {
142 label: Some("Trueno GPU Device"),
143 required_features: wgpu::Features::empty(),
144 required_limits: limits,
145 memory_hints: wgpu::MemoryHints::Performance,
146 experimental_features: Default::default(),
147 trace: Default::default(),
148 })
149 .await
150 .map_err(|e| format!("Failed to create device: {}", e))?;
151
152 Ok(Self { device, queue })
153 }
154
155 #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
160 pub fn new_with_adapter_index(index: u32) -> Result<Self, String> {
161 let _guard = DEVICE_INIT_LOCK.lock().unwrap_or_else(|e| e.into_inner());
163 runtime::block_on(async { Self::new_with_adapter_index_async(index).await })
164 }
165
166 pub async fn new_with_adapter_index_async(index: u32) -> Result<Self, String> {
171 let instance = shared_instance();
172 let adapters = instance.enumerate_adapters(gpu_backends());
174
175 if adapters.is_empty() {
176 return Err("No GPU adapters found".to_string());
177 }
178
179 let adapter = adapters
180 .into_iter()
181 .nth(index as usize)
182 .ok_or_else(|| format!("GPU adapter index {} out of range", index))?;
183
184 let mut limits = wgpu::Limits::default();
185 limits.max_buffer_size = adapter.limits().max_buffer_size;
186 limits.max_storage_buffer_binding_size = adapter.limits().max_storage_buffer_binding_size;
187
188 let (device, queue) = adapter
189 .request_device(&wgpu::DeviceDescriptor {
190 label: Some(&format!("Trueno GPU Device [{}]", index)),
191 required_features: wgpu::Features::empty(),
192 required_limits: limits,
193 memory_hints: wgpu::MemoryHints::Performance,
194 experimental_features: Default::default(),
195 trace: Default::default(),
196 })
197 .await
198 .map_err(|e| format!("Failed to create device at index {}: {}", index, e))?;
199
200 Ok(Self { device, queue })
201 }
202
203 #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
207 pub fn list_adapters() -> Vec<(u32, String, String)> {
208 let _guard = DEVICE_INIT_LOCK.lock().unwrap_or_else(|e| e.into_inner());
211 runtime::block_on(Self::list_adapters_async())
212 }
213
214 pub async fn list_adapters_async() -> Vec<(u32, String, String)> {
216 let instance = shared_instance();
217 let adapters = instance.enumerate_adapters(gpu_backends());
219
220 adapters
221 .iter()
222 .enumerate()
223 .map(|(idx, adapter)| {
224 let info = adapter.get_info();
225 (idx as u32, info.name, format!("{:?}", info.backend))
226 })
227 .collect()
228 }
229
230 #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
247 pub fn is_available() -> bool {
248 use std::sync::OnceLock;
249 static AVAILABLE: OnceLock<bool> = OnceLock::new();
250 *AVAILABLE.get_or_init(|| runtime::block_on(Self::is_available_async()))
251 }
252
253 pub async fn is_available_async() -> bool {
255 let instance = shared_instance();
256 instance
257 .request_adapter(&wgpu::RequestAdapterOptions {
258 power_preference: wgpu::PowerPreference::HighPerformance,
259 compatible_surface: None,
260 force_fallback_adapter: false,
261 })
262 .await
263 .is_ok()
264 }
265
266 pub(super) async fn execute_element_wise_op(
279 &self,
280 op_name: &str,
281 shader_source: &str,
282 input: &[f32],
283 result: &mut [f32],
284 uniform_data: Option<&[u8]>,
285 ) -> Result<(), String> {
286 let len = input.len();
287
288 let shader = self.device.create_shader_module(wgpu::ShaderModuleDescriptor {
290 label: Some(&format!("{} Shader", op_name)),
291 source: wgpu::ShaderSource::Wgsl(shader_source.into()),
292 });
293
294 let input_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
296 label: Some(&format!("{} Input", op_name)),
297 size: std::mem::size_of_val(input) as u64,
298 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
299 mapped_at_creation: false,
300 });
301
302 let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
304 label: Some(&format!("{} Output", op_name)),
305 size: std::mem::size_of_val(result) as u64,
306 usage: wgpu::BufferUsages::STORAGE
307 | wgpu::BufferUsages::COPY_SRC
308 | wgpu::BufferUsages::COPY_DST,
309 mapped_at_creation: false,
310 });
311
312 self.queue.write_buffer(&input_buffer, 0, bytemuck::cast_slice(input));
314
315 let uniform_buffer = uniform_data.map(|data| {
317 let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
318 label: Some(&format!("{} Uniform", op_name)),
319 size: data.len() as u64,
320 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
321 mapped_at_creation: false,
322 });
323 self.queue.write_buffer(&buffer, 0, data);
324 buffer
325 });
326
327 let mut bind_group_entries = vec![
329 wgpu::BindGroupLayoutEntry {
330 binding: 0,
331 visibility: wgpu::ShaderStages::COMPUTE,
332 ty: wgpu::BindingType::Buffer {
333 ty: wgpu::BufferBindingType::Storage { read_only: true },
334 has_dynamic_offset: false,
335 min_binding_size: None,
336 },
337 count: None,
338 },
339 wgpu::BindGroupLayoutEntry {
340 binding: 1,
341 visibility: wgpu::ShaderStages::COMPUTE,
342 ty: wgpu::BindingType::Buffer {
343 ty: wgpu::BufferBindingType::Storage { read_only: false },
344 has_dynamic_offset: false,
345 min_binding_size: None,
346 },
347 count: None,
348 },
349 ];
350
351 if uniform_buffer.is_some() {
353 bind_group_entries.push(wgpu::BindGroupLayoutEntry {
354 binding: 2,
355 visibility: wgpu::ShaderStages::COMPUTE,
356 ty: wgpu::BindingType::Buffer {
357 ty: wgpu::BufferBindingType::Uniform,
358 has_dynamic_offset: false,
359 min_binding_size: None,
360 },
361 count: None,
362 });
363 }
364
365 let bind_group_layout =
367 self.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
368 label: Some(&format!("{} Bind Group Layout", op_name)),
369 entries: &bind_group_entries,
370 });
371
372 let mut bind_entries = vec![
374 wgpu::BindGroupEntry { binding: 0, resource: input_buffer.as_entire_binding() },
375 wgpu::BindGroupEntry { binding: 1, resource: output_buffer.as_entire_binding() },
376 ];
377
378 if let Some(ref uniform_buf) = uniform_buffer {
380 bind_entries.push(wgpu::BindGroupEntry {
381 binding: 2,
382 resource: uniform_buf.as_entire_binding(),
383 });
384 }
385
386 let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
388 label: Some(&format!("{} Bind Group", op_name)),
389 layout: &bind_group_layout,
390 entries: &bind_entries,
391 });
392
393 let pipeline_layout = self.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
395 label: Some(&format!("{} Pipeline Layout", op_name)),
396 bind_group_layouts: &[&bind_group_layout],
397 push_constant_ranges: &[],
398 });
399
400 let pipeline = self.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
401 label: Some(&format!("{} Pipeline", op_name)),
402 layout: Some(&pipeline_layout),
403 module: &shader,
404 entry_point: Some("main"),
405 compilation_options: Default::default(),
406 cache: None,
407 });
408
409 let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
411 label: Some(&format!("{} Staging Buffer", op_name)),
412 size: std::mem::size_of_val(result) as u64,
413 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
414 mapped_at_creation: false,
415 });
416
417 let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
419 label: Some(&format!("{} Encoder", op_name)),
420 });
421
422 {
423 let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
424 label: Some(&format!("{} Pass", op_name)),
425 timestamp_writes: None,
426 });
427 compute_pass.set_pipeline(&pipeline);
428 compute_pass.set_bind_group(0, &bind_group, &[]);
429
430 let workgroup_size = 256;
432 let num_workgroups = (len as u32).div_ceil(workgroup_size);
433
434 compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
435 }
436
437 encoder.copy_buffer_to_buffer(
439 &output_buffer,
440 0,
441 &staging_buffer,
442 0,
443 std::mem::size_of_val(result) as u64,
444 );
445
446 self.queue.submit(Some(encoder.finish()));
448
449 let buffer_slice = staging_buffer.slice(..);
451 let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
452 buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
453 sender.send(result).ok();
454 });
455
456 self.device.poll(wgpu::PollType::Wait { submission_index: None, timeout: None }).ok();
458
459 receiver
460 .receive()
461 .await
462 .ok_or("Failed to receive mapping result")?
463 .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
464
465 {
466 let data = buffer_slice.get_mapped_range();
467 result.copy_from_slice(bytemuck::cast_slice(&data));
468 }
469
470 staging_buffer.unmap();
471
472 Ok(())
473 }
474}
475
476#[cfg(all(test, feature = "gpu", not(target_arch = "wasm32")))]
477mod tests {
478 use super::*;
479
480 #[test]
487 fn test_gpu_backends_excludes_gles() {
488 let mask = gpu_backends();
489
490 assert!(
492 !mask.contains(wgpu::Backends::GL),
493 "gpu_backends() must NOT include Backends::GL (GLES/EGL panics in Drop \
494 on Linux/AMD-RADV → SIGABRT). mask = {:?}",
495 mask
496 );
497
498 #[cfg(any(target_os = "linux", target_os = "android"))]
500 assert!(
501 mask.contains(wgpu::Backends::VULKAN),
502 "gpu_backends() must include VULKAN on Linux (AMD-RADV/NVIDIA). mask = {:?}",
503 mask
504 );
505 #[cfg(target_os = "macos")]
506 assert!(
507 mask.contains(wgpu::Backends::METAL),
508 "gpu_backends() must include METAL on macOS (Apple Silicon). mask = {:?}",
509 mask
510 );
511 #[cfg(target_os = "windows")]
512 assert!(
513 mask.contains(wgpu::Backends::VULKAN) || mask.contains(wgpu::Backends::DX12),
514 "gpu_backends() must include VULKAN or DX12 on Windows. mask = {:?}",
515 mask
516 );
517 }
518
519 #[test]
520 fn test_is_available_consistency() {
521 let available = GpuDevice::is_available();
524 let device_result = GpuDevice::new();
525
526 if available {
527 assert!(
529 device_result.is_ok(),
530 "is_available() returned true, but GpuDevice::new() failed"
531 );
532 } else {
533 eprintln!(
537 "GPU not available (is_available=false), device creation result: {:?}",
538 device_result.is_err()
539 );
540 }
541 }
542
543 #[test]
544 fn test_reduce_sum_not_hardcoded() {
545 if !GpuDevice::is_available() {
547 eprintln!("GPU not available, skipping test");
548 return;
549 }
550
551 let device = GpuDevice::new().expect("Failed to create GPU device");
552 let input = vec![1.0, 2.0, 3.0, 4.0, 5.0]; let result = runtime::block_on(device.reduce_sum(&input)).expect("reduce_sum failed");
556
557 assert_ne!(result, -1.0, "reduce_sum returned hardcoded -1.0 (mutant not killed)");
559
560 let expected: f32 = input.iter().sum();
562 assert!(
563 (result - expected).abs() < 1e-4,
564 "reduce_sum({:?}) = {} (expected {})",
565 input,
566 result,
567 expected
568 );
569 }
570}
571
572#[cfg(all(test, feature = "gpu", not(target_arch = "wasm32")))]
579mod pmat952_tests {
580 use super::{shared_instance, DEVICE_INIT_LOCK};
581
582 const INNER_ENV: &str = "PMAT952_INNER";
583
584 #[test]
585 fn test_pmat952_shared_instance_does_not_deadlock_against_device_init_lock() {
586 if std::env::var_os(INNER_ENV).is_some() {
587 return; }
589 let exe = std::env::current_exe().expect("current_exe");
590 let mut child = std::process::Command::new(exe)
591 .args([
592 "--exact",
593 "backends::gpu::device::pmat952_tests::pmat952_inner_probe",
594 "--nocapture",
595 "--test-threads=1",
596 ])
597 .env(INNER_ENV, "1")
598 .stdout(std::process::Stdio::null())
599 .stderr(std::process::Stdio::null())
600 .spawn()
601 .expect("spawn the probe child");
602 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(90);
603 loop {
604 match child.try_wait().expect("try_wait") {
605 Some(status) => {
606 assert!(
607 status.success(),
608 "PMAT-952: the probe child failed ({status}): shared_instance() parked \
609 behind DEVICE_INIT_LOCK in a fresh process (the ABBA deadlock)"
610 );
611 return;
612 }
613 None if std::time::Instant::now() > deadline => {
614 let _ = child.kill();
615 panic!("PMAT-952: the probe child hung: shared_instance() parked behind DEVICE_INIT_LOCK");
616 }
617 None => std::thread::sleep(std::time::Duration::from_millis(100)),
618 }
619 }
620 }
621
622 #[test]
624 fn pmat952_inner_probe() {
625 if std::env::var_os(INNER_ENV).is_none() {
626 return;
627 }
628 let _guard = DEVICE_INIT_LOCK.lock().unwrap_or_else(|e| e.into_inner());
629 let (tx, rx) = std::sync::mpsc::channel();
630 std::thread::spawn(move || {
631 let _instance = shared_instance();
632 let _ = tx.send(());
633 });
634 rx.recv_timeout(std::time::Duration::from_secs(30))
635 .expect("shared_instance() parked while DEVICE_INIT_LOCK was held (PMAT-952)");
636 }
637}