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/// Re-export wgpu types for downstream crates that need to create persistent
70/// GPU buffers (KAIZEN-015: GPU-resident weights).
71#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
72pub use wgpu;
73
74#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
75pub use pool::GpuDevicePool;
76
77/// PMAT-322: Cached matmul with persistent weight buffers for LLM inference.
78#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
79pub use device::linalg::cached_matmul::GpuMatmulCache;
80
81/// PMAT-324: WGSL transformer forward pass shaders.
82#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
83pub use device::linalg::wgsl_forward::{QkvLoRA, WgslForwardPass};
84
85#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
86mod backend_ops;
87
88/// GPU backend for compute operations (native only, uses sync wrappers)
89#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
90#[derive(Clone)]
91pub struct GpuBackend {
92    device: Option<GpuDevice>,
93}
94
95#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
96impl GpuBackend {
97    /// Create a new GPU backend
98    pub fn new() -> Self {
99        Self { device: None }
100    }
101
102    /// Initialize GPU device (lazy)
103    fn ensure_device(&mut self) -> Result<&GpuDevice, String> {
104        if self.device.is_none() {
105            self.device = Some(GpuDevice::new()?);
106        }
107        Ok(self.device.as_ref().expect("device initialized above"))
108    }
109
110    /// Check if GPU is available
111    pub fn is_available() -> bool {
112        GpuDevice::is_available()
113    }
114}
115
116#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
117impl Default for GpuBackend {
118    fn default() -> Self {
119        Self::new()
120    }
121}
122
123// Stub implementation when GPU feature is disabled or on WASM
124#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
125#[derive(Clone)]
126pub struct GpuBackend;
127
128#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
129impl GpuBackend {
130    pub fn new() -> Self {
131        Self
132    }
133
134    pub fn is_available() -> bool {
135        false
136    }
137}
138
139#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
140impl Default for GpuBackend {
141    fn default() -> Self {
142        Self
143    }
144}
145
146// Tests for stub implementation (when GPU feature is NOT enabled)
147#[cfg(test)]
148#[cfg(not(feature = "gpu"))]
149mod stub_tests {
150    use super::*;
151
152    #[test]
153    fn test_gpu_backend_stub_new() {
154        let _backend = GpuBackend::new();
155    }
156
157    #[test]
158    fn test_gpu_backend_stub_is_available() {
159        assert!(!GpuBackend::is_available());
160    }
161
162    #[test]
163    fn test_gpu_backend_stub_default() {
164        let _ = GpuBackend;
165    }
166
167    #[test]
168    fn test_gpu_backend_stub_clone() {
169        let backend = GpuBackend::new();
170        let _cloned = backend.clone();
171    }
172}
173
174#[cfg(test)]
175#[cfg(feature = "gpu")]
176mod tests_gpu;