Skip to main content

trueno/backends/gpu/
mod.rs

1#![allow(missing_docs)]
2//! GPU backend using wgpu (Vulkan/Metal/DX12/WebGPU)
3//!
4//! This backend provides GPU-accelerated compute for large-scale operations.
5//! It uses wgpu for cross-platform GPU access and WGSL compute shaders.
6//!
7//! # Performance
8//!
9//! GPU backend is optimal for very large workloads (>100K elements for reductions,
10//! >1000×1000 for matrix operations) where transfer overhead is amortized.
11//!
12//! Expected speedups vs SIMD:
13//! - Matrix multiplication (large): 10-50x
14//! - Reductions (large): 5-20x
15//!
16//! # Architecture
17//!
18//! - Device initialization is lazy (first GPU operation)
19//! - Compute shaders written in WGSL
20//! - Asynchronous execution with pollster for blocking
21//! - Automatic fallback to CPU if GPU unavailable
22//!
23//! # Memory Hierarchy Abstractions
24//!
25//! - [`TensorView`] - Structured view into GPU memory with shape/stride metadata
26//! - [`PartitionView`] - Tiling strategy for efficient GPU work distribution
27//!
28//! Based on cuda-tile-behavior.md Section 3.2.
29
30#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
31mod batch;
32
33#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
34mod device;
35
36#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
37mod pool;
38
39#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
40pub mod shaders;
41
42#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
43pub mod runtime;
44
45// Memory hierarchy abstractions (always available, no GPU feature required)
46mod partition_view;
47mod tensor_view;
48mod tiled_reduction;
49
50pub use partition_view::{PartitionView, TileInfo};
51pub use tensor_view::{MemoryLayout, TensorView};
52pub use tiled_reduction::{
53    tiled_max_2d, tiled_min_2d, tiled_reduce_2d, tiled_reduce_partial, tiled_sum_2d, MaxOp, MinOp,
54    ReduceOp, SumOp, TILE_SIZE,
55};
56
57#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
58pub use batch::{BufferId, GpuCommandBatch, PipelineCache};
59
60// Export GpuDevice for both native and WASM GPU features
61#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
62pub use device::GpuDevice;
63
64// PMAT-778: process-global shared wgpu instance, reused by every crate-internal
65// adapter/device enumeration so the broken freedreno ICD (GB10) is touched once.
66#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
67pub(crate) use device::shared_instance;
68
69// PMAT-925: non-GLES backend mask, reused by every crate-internal
70// `enumerate_adapters` call so the broken GLES/EGL adapter (SIGABRT-in-Drop on
71// Linux/AMD-RADV) is never instantiated.
72#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
73pub(crate) use device::gpu_backends;
74
75/// Re-export wgpu types for downstream crates that need to create persistent
76/// GPU buffers (KAIZEN-015: GPU-resident weights).
77#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
78pub use wgpu;
79
80#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
81pub use pool::GpuDevicePool;
82
83/// PMAT-322: Cached matmul with persistent weight buffers for LLM inference.
84#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
85pub use device::linalg::cached_matmul::GpuMatmulCache;
86
87/// PMAT-324: WGSL transformer forward pass shaders.
88#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
89pub use device::linalg::wgsl_forward::{QkvLoRA, WgslForwardPass};
90
91#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
92mod backend_ops;
93
94/// GPU backend for compute operations (native only, uses sync wrappers)
95#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
96#[derive(Clone)]
97pub struct GpuBackend {
98    device: Option<GpuDevice>,
99}
100
101#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
102impl GpuBackend {
103    /// Create a new GPU backend
104    pub fn new() -> Self {
105        Self { device: None }
106    }
107
108    /// Initialize GPU device (lazy)
109    fn ensure_device(&mut self) -> Result<&GpuDevice, String> {
110        if self.device.is_none() {
111            self.device = Some(GpuDevice::new()?);
112        }
113        Ok(self.device.as_ref().expect("device initialized above"))
114    }
115
116    /// Check if GPU is available
117    pub fn is_available() -> bool {
118        GpuDevice::is_available()
119    }
120}
121
122#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
123impl Default for GpuBackend {
124    fn default() -> Self {
125        Self::new()
126    }
127}
128
129// Stub implementation when GPU feature is disabled or on WASM
130#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
131#[derive(Clone)]
132pub struct GpuBackend;
133
134#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
135impl GpuBackend {
136    pub fn new() -> Self {
137        Self
138    }
139
140    pub fn is_available() -> bool {
141        false
142    }
143}
144
145#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
146impl Default for GpuBackend {
147    fn default() -> Self {
148        Self
149    }
150}
151
152// Tests for stub implementation (when GPU feature is NOT enabled)
153#[cfg(test)]
154#[cfg(not(feature = "gpu"))]
155mod stub_tests {
156    use super::*;
157
158    #[test]
159    fn test_gpu_backend_stub_new() {
160        let _backend = GpuBackend::new();
161    }
162
163    #[test]
164    fn test_gpu_backend_stub_is_available() {
165        assert!(!GpuBackend::is_available());
166    }
167
168    #[test]
169    fn test_gpu_backend_stub_default() {
170        let _ = GpuBackend;
171    }
172
173    #[test]
174    fn test_gpu_backend_stub_clone() {
175        let backend = GpuBackend::new();
176        let _cloned = backend.clone();
177    }
178}
179
180#[cfg(test)]
181#[cfg(feature = "gpu")]
182mod tests_gpu;