Skip to main content

trueno/backends/gpu/device/
mod.rs

1//! GPU device initialization and management
2//!
3//! This module provides cross-platform GPU compute via wgpu (WebGPU).
4//!
5//! # Platform differences
6//!
7//! - **Native**: Sync wrappers available using `pollster::block_on`
8//! - **WASM**: Sync wrappers unavailable (can't block main thread); use `*_async` methods
9//!
10//! Use `runtime::sync_available()` to check at runtime.
11
12mod 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/// Process-global lock serializing native wgpu instance/adapter/device creation.
22///
23/// PMAT-778: On hosts whose Vulkan ICD is unsafe to initialize concurrently
24/// (notably the NVIDIA GB10 / aarch64, where the Mesa freedreno "Turnip" ICD
25/// probes `/dev/dri/renderD128` and **segfaults** when multiple threads create
26/// `wgpu::Instance`s and request adapters/devices simultaneously), every sync
27/// device-creation entry point takes this lock. Device creation is a rare,
28/// one-time-per-scheduler operation off the compute hot path, so serializing it
29/// is free on healthy GPUs and turns a concurrent-init crash into ordered,
30/// correct initialization. This is the companion to the adapter-probe `OnceLock`
31/// cache (PMAT-773): the probe is memoized once, and actual device acquisition
32/// is serialized here.
33#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
34static DEVICE_INIT_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
35
36/// Platform-appropriate wgpu backend mask for adapter enumeration.
37///
38/// PMAT-925: `wgpu::Backends::all()` includes [`wgpu::Backends::GL`], which on
39/// Linux hosts that have both Vulkan and GLES/EGL (notably the intel AMD-RADV
40/// cross-silicon baseline box) instantiates a GLES adapter whose
41/// `EglContext::make_current` **panics inside `Drop`** (wgpu-hal-27.0.4
42/// `gles/egl.rs:305`). A panic in a destructor during cleanup aborts the whole
43/// process with SIGABRT ("panic in a destructor during cleanup"); standalone
44/// `list_adapters()` / `is_available()` could also spin/hang on the broken EGL
45/// path. The compute kernels themselves are correct — the fragility is purely in
46/// adapter *enumeration* and the GLES `Drop` path.
47///
48/// We return [`wgpu::Backends::PRIMARY`], which in wgpu 27 is
49/// `VULKAN | METAL | DX12 | BROWSER_WEBGPU` and **excludes** `GL` (GL lives only
50/// in `Backends::SECONDARY`). This keeps the real GPU on every platform — Vulkan
51/// on Linux/AMD-RADV, Metal on Apple, DX12 on Windows — while guaranteeing the
52/// broken GLES/EGL adapter is never created.
53///
54/// The mask is applied at BOTH the `wgpu::Instance` construction site (so the
55/// GLES backend is never even registered on the instance) and every
56/// `enumerate_adapters` call site.
57#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
58pub(crate) const fn gpu_backends() -> wgpu::Backends {
59    // PRIMARY = VULKAN | METAL | DX12 | BROWSER_WEBGPU (never GL/GLES).
60    wgpu::Backends::PRIMARY
61}
62
63/// Process-global, lazily-created shared `wgpu::Instance`.
64///
65/// PMAT-778: Creating a fresh `wgpu::Instance` enumerates every installed Vulkan
66/// ICD. On the NVIDIA GB10 / aarch64 the host ships ~11 Mesa ICDs (freedreno
67/// "Turnip", panfrost, asahi, …) that are irrelevant to this hardware; the
68/// freedreno ICD spawns its own background Vulkan threads and **segfaults** when
69/// several instances enumerate it concurrently (each open of
70/// `/dev/dri/renderD128` returns `VK_ERROR_INCOMPATIBLE_DRIVER`). Sharing one
71/// instance for the whole process means the broken ICD is enumerated exactly
72/// once, eliminating the concurrent-init race entirely. `wgpu::Instance` is
73/// `Clone`/`Send`/`Sync`, so every adapter/device request can cheaply reuse it.
74#[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            // PMAT-952: this initializer must NOT take DEVICE_INIT_LOCK. The
81            // `OnceLock` already serializes the one-time enumeration (every
82            // other caller parks on it until this closure returns), and
83            // `GpuDevice::new` holds DEVICE_INIT_LOCK while it calls
84            // `shared_instance()`. Taking the mutex here as well was an ABBA
85            // cycle: a thread entering through `is_available_async` (no mutex)
86            // became the initializer, blocked on the mutex held by a
87            // `GpuDevice::new` thread, which was parked on this `OnceLock`.
88            // Measured 2026-09-04: 49 threads parked, release clean-room hung.
89            // PMAT-925: constrain the instance to a non-GLES backend mask so the
90            // broken GLES/EGL adapter (SIGABRT-in-Drop on Linux/AMD-RADV) is never
91            // registered. `Instance::default()` would use `Backends::all()`
92            // (which includes GL).
93            wgpu::Instance::new(&wgpu::InstanceDescriptor {
94                backends: gpu_backends(),
95                ..Default::default()
96            })
97        })
98        .clone()
99}
100
101/// GPU device manager
102#[derive(Clone)]
103pub struct GpuDevice {
104    pub device: wgpu::Device,
105    pub queue: wgpu::Queue,
106}
107
108impl GpuDevice {
109    /// Initialize GPU device (sync, native only)
110    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
111    pub fn new() -> Result<Self, String> {
112        // PMAT-778: serialize concurrent native device creation (freedreno ICD
113        // segfaults on the GB10 when instances/devices are created in parallel).
114        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    /// Initialize GPU device (async, works on all platforms)
119    pub async fn new_async() -> Result<Self, String> {
120        // Create instance
121        let instance = shared_instance();
122
123        // Request adapter (GPU)
124        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        // Request device and queue with adapter's actual max buffer size
134        // Default wgpu limits cap buffers at 256MB, which is too small for
135        // 7B+ model weight matrices (e.g., FFN [18944, 3584] x f32 = 271MB)
136        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    /// Initialize GPU device with a specific adapter index (sync, native only)
156    ///
157    /// Use this to select a specific GPU when multiple are available.
158    /// Adapter indices correspond to `Instance::enumerate_adapters()` ordering.
159    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
160    pub fn new_with_adapter_index(index: u32) -> Result<Self, String> {
161        // PMAT-778: serialize concurrent native device creation (see DEVICE_INIT_LOCK).
162        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    /// Initialize GPU device with a specific adapter index (async, all platforms)
167    ///
168    /// Use this to select a specific GPU when multiple are available.
169    /// Adapter indices correspond to `Instance::enumerate_adapters()` ordering.
170    pub async fn new_with_adapter_index_async(index: u32) -> Result<Self, String> {
171        let instance = shared_instance();
172        // PMAT-925: exclude GLES (see `gpu_backends`).
173        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    /// List all available GPU adapters (sync, native only)
204    ///
205    /// Returns a list of (index, name, backend) tuples for each adapter.
206    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
207    pub fn list_adapters() -> Vec<(u32, String, String)> {
208        // PMAT-778: serialize concurrent native instance/adapter enumeration
209        // (freedreno ICD segfaults on the GB10 under concurrent init).
210        let _guard = DEVICE_INIT_LOCK.lock().unwrap_or_else(|e| e.into_inner());
211        runtime::block_on(Self::list_adapters_async())
212    }
213
214    /// List all available GPU adapters (async, all platforms)
215    pub async fn list_adapters_async() -> Vec<(u32, String, String)> {
216        let instance = shared_instance();
217        // PMAT-925: exclude GLES (see `gpu_backends`).
218        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    /// Check if GPU is available (sync, native only)
231    ///
232    /// PMAT-773: The adapter probe is cached for the lifetime of the process via a
233    /// [`std::sync::OnceLock`]. On hosts WITHOUT a wgpu-compatible adapter (e.g.
234    /// headless NVIDIA / Jetson / Blackwell GB10, where opening
235    /// `/dev/dri/renderD128` fails with `VK_ERROR_INCOMPATIBLE_DRIVER`), the
236    /// underlying `request_adapter` call is expensive and was being re-attempted by
237    /// every test/caller, adding diffuse latency across a test suite. Caching the
238    /// result means the failing Vulkan device open is attempted at most once per
239    /// process; subsequent callers short-circuit.
240    ///
241    /// Behavior on hosts WITH a usable GPU is unchanged: the first probe succeeds,
242    /// `true` is cached, and callers proceed to acquire a real device via
243    /// [`Self::new`] exactly as before. The cache only memoizes whether an adapter
244    /// is obtainable — it never holds a device, so real-GPU acquisition is never
245    /// short-circuited.
246    #[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    /// Check if GPU is available (async, works on all platforms)
254    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    /// Generic helper for element-wise GPU operations
267    ///
268    /// This helper eliminates code duplication between element-wise operations
269    /// (relu, clip, sigmoid, tanh, etc.) by abstracting the common GPU compute pattern.
270    ///
271    /// # Arguments
272    ///
273    /// * `op_name` - Operation name for labels (e.g., "ReLU", "Clip")
274    /// * `shader_source` - WGSL shader source code
275    /// * `input` - Input data
276    /// * `result` - Output buffer
277    /// * `uniform_data` - Optional uniform buffer data (e.g., clip parameters)
278    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        // Create shader module
289        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        // Create input buffer
295        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        // Create output buffer
303        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        // Write input data
313        self.queue.write_buffer(&input_buffer, 0, bytemuck::cast_slice(input));
314
315        // Create optional uniform buffer
316        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        // Create bind group layout entries (input + output + optional uniform)
328        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        // Add uniform buffer binding if present
352        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        // Create bind group layout
366        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        // Create bind group entries
373        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        // Add uniform buffer binding if present
379        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        // Create bind group
387        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        // Create pipeline
394        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        // Create staging buffer for reading results
410        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        // Create command encoder
418        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            // Dispatch workgroups (256 threads per workgroup)
431            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        // Copy result to staging buffer
438        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        // Submit commands
447        self.queue.submit(Some(encoder.finish()));
448
449        // Read back results
450        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        // Poll device to ensure GPU work completes and callbacks are invoked
457        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    /// PMAT-925 FALSIFIER: the adapter-enumeration backend mask MUST NOT contain
481    /// GLES (`wgpu::Backends::GL`), and MUST contain the platform's real backend.
482    ///
483    /// RED on `Backends::all()` (contains GL → GLES/EGL adapter → SIGABRT-in-Drop
484    /// on Linux/AMD-RADV). GREEN on `Backends::PRIMARY`. Host-independent: it
485    /// inspects the bitmask, it does not create any adapter.
486    #[test]
487    fn test_gpu_backends_excludes_gles() {
488        let mask = gpu_backends();
489
490        // The whole point: GLES/EGL must never be enumerated.
491        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        // The real GPU backend on each platform must still be present.
499        #[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        // EXTREME TDD: Kill mutant that replaces is_available() with hardcoded false
522        // Test that is_available() is consistent with GpuDevice::new()
523        let available = GpuDevice::is_available();
524        let device_result = GpuDevice::new();
525
526        if available {
527            // If is_available() returns true, device creation should succeed
528            assert!(
529                device_result.is_ok(),
530                "is_available() returned true, but GpuDevice::new() failed"
531            );
532        } else {
533            // If is_available() returns false, we can't make assertions about new()
534            // (it might still succeed in some edge cases, but typically should fail)
535            // The key test is: mutant always returns false, so on GPU systems this fails
536            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        // EXTREME TDD: Kill mutant that replaces reduce_sum with Ok(-1.0)
546        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]; // sum = 15.0
553
554        // reduce_sum is async, so we use runtime::block_on
555        let result = runtime::block_on(device.reduce_sum(&input)).expect("reduce_sum failed");
556
557        // Kill mutant: verify result is NOT -1.0
558        assert_ne!(result, -1.0, "reduce_sum returned hardcoded -1.0 (mutant not killed)");
559
560        // Verify correct computation
561        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/// PMAT-952: the `shared_instance` initializer must never take
573/// `DEVICE_INIT_LOCK`. `GpuDevice::new` holds that mutex while it calls
574/// `shared_instance()`, so an initializer that also took it was an ABBA cycle:
575/// the release clean-room of 2026-09-04 parked 49 threads on it. The race is
576/// only reachable on the FIRST initialization of a process, so the probe runs
577/// in a fresh child process where the `OnceLock` is still empty.
578#[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; // we are the child; the probe below is the subject
588        }
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    /// Runs only inside the child spawned above (fresh process, empty `OnceLock`).
623    #[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}