optirs_gpu/
lib.rs

1//! # OptiRS GPU — GPU acceleration and GPU-aware optimizer tooling
2//!
3//! **Version:** 0.3.2
4//!
5//! `optirs-gpu` has two halves, and it is worth being precise about which is
6//! which:
7//!
8//! 1. **A real GPU optimizer path.** [`optimizers`] runs Adam, AdamW, SGD,
9//!    RMSprop, Adagrad and LAMB as compute shaders through
10//!    [`scirs2_core::gpu`]. Parameters and gradients are uploaded to device
11//!    buffers, a compiled pipeline is dispatched, and the result is read back;
12//!    the per-parameter optimizer state stays resident in device memory between
13//!    steps. The kernels ship in both WGSL and MSL ([`shaders`]) so the same
14//!    optimizer runs on whichever backend the machine can reach.
15//! 2. **A CPU library of GPU-*aware* algorithms.** [`occupancy`],
16//!    [`kernel_fusion`], [`quantization`], [`sparse_optimizer`] and
17//!    [`memory::allocation`] / [`memory::management`] are pure-CPU models,
18//!    planners and numerical routines that reason *about* GPU execution. They
19//!    have no device dependency and are fully covered by unit tests.
20//!
21//! ## Backend support matrix
22//!
23//! | Backend | Status |
24//! |---------|--------|
25//! | Metal (`metal`, automatic on macOS) | ✅ real compute: MSL pipelines, buffers, dispatch, readback |
26//! | WebGPU (`wgpu`, default) | ✅ WGSL kernels are implemented, but `scirs2-core` 0.6.5's runtime device probe never enumerates wgpu adapters, so `GpuContext::new(Wgpu)` currently fails everywhere. The path goes live when that probe is fixed |
27//! | OpenCL (`opencl`) | 🚧 context creation only — no OpenCL C kernel sources are shipped |
28//! | CUDA (`cuda`) | ❌ not available — `scirs2-core` removed its CUDA backend in 0.6.x; the feature gates reporting code only |
29//! | ROCm | ❌ not available |
30//!
31//! [`optimizers::GpuOptimizerConfig`] defaults to probing
32//! [`optimizers::SUPPORTED_BACKENDS`] in order and using the first that opens.
33//!
34//! ## Not implemented (and not faked)
35//!
36//! * Cross-**device** collectives. [`multi_gpu`] can drive a real reduction
37//!   kernel on a single device; anything that would require moving data
38//!   between two physical GPUs returns
39//!   [`GpuOptimError::UnsupportedOperation`].
40//! * Literal NVIDIA tensor cores / `wmma`. [`tensor_cores`] provides real
41//!   CPU-side matrix-layout optimization, precision selection and AMP loss
42//!   scaling; the device GEMM entry points (`tensor_core_gemm`,
43//!   `fused_adam_tensor_core`, ...) report an honest
44//!   [`GpuOptimError::UnsupportedOperation`] because no backend this crate can
45//!   reach exposes NVIDIA tensor cores.
46//!
47//! ## Example
48//!
49//! ```no_run
50//! use optirs_gpu::optimizers::{AdamParams, GpuAdam};
51//! use optirs_gpu::GpuOptimizer;
52//! use scirs2_core::ndarray::Array1;
53//!
54//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
55//! let mut optimizer = GpuAdam::new(AdamParams::default())?;
56//! optimizer.move_to_gpu()?;
57//!
58//! let mut params = Array1::from_elem(1_024, 1.0f32);
59//! let grads = Array1::from_elem(1_024, 0.01f32);
60//! optimizer.step_gpu(&mut params, &grads)?;
61//!
62//! // Bring the moment estimates back to host memory when done.
63//! optimizer.move_to_cpu()?;
64//! # Ok(())
65//! # }
66//! ```
67//!
68//! ## Architecture
69//!
70//! Every device access goes through SciRS2:
71//! - **GPU context**: `scirs2_core::gpu::GpuContext`
72//! - **GPU memory**: `scirs2_core::gpu::GpuBuffer`
73//! - **Kernel compilation**: `scirs2_core::gpu::GpuCompiler`
74//!
75//! `wgpu`, `pollster`, `metal` and friends are *not* direct dependencies of
76//! this crate; they arrive through the `scirs2-core/<backend>` features that
77//! this crate's features forward.
78
79use scirs2_core::gpu::GpuError;
80use scirs2_core::ndarray::{Array, Dimension};
81use scirs2_core::numeric::Float;
82
83pub mod backends;
84pub mod kernel_fusion;
85pub mod memory;
86pub mod mixed_precision;
87pub mod multi_gpu;
88pub mod occupancy;
89pub mod optimizers;
90pub mod quantization;
91pub mod shaders;
92pub mod sparse_optimizer;
93pub mod tensor_cores;
94pub mod utils;
95
96pub use backends::GpuBackend;
97pub use kernel_fusion::{FusionGraph, FusionGroup, FusionOp, FusionPlan, FusionPlanner, OpKind};
98pub use memory::MemoryPool;
99pub use mixed_precision::{
100    f16_bits_to_f32, f32_to_f16_bits, DynamicLossScaler, MixedPrecisionConfig, OverflowStats,
101};
102pub use occupancy::{
103    calculate_occupancy, optimal_block_size, KernelResourceUsage, OccupancyLimiter,
104    OccupancyResult, SmResourceLimits,
105};
106pub use optimizers::{
107    AdagradParams, AdamParams, GpuAdagrad, GpuAdam, GpuAdamW, GpuLamb, GpuOptimizerConfig,
108    GpuRmsprop, GpuSgd, RmspropParams, SgdParams,
109};
110pub use quantization::{
111    fake_quant_backward, fake_quant_fp8, fake_quant_int, fake_quant_int_per_channel,
112    per_channel_params, Fp8Format, IntDtype, QatConfig, QatOptimizer, QuantParams, QuantScheme,
113    QuantTarget, RoundingMode,
114};
115pub use sparse_optimizer::{
116    CooGradient, CsrGradient, LazyAdamMode, SparseAdam, SparseAdamConfig, SparseAdamTable,
117    SparseSgd, SparseSgdConfig, SparseSgdTable,
118};
119
120/// Error type for GPU optimizer operations
121#[derive(Debug, thiserror::Error)]
122pub enum GpuOptimError {
123    /// GPU backend error
124    #[error("GPU error: {0}")]
125    GpuError(#[from] GpuError),
126
127    /// Unsupported operation
128    #[error("Operation not supported: {0}")]
129    UnsupportedOperation(String),
130
131    /// Invalid state
132    #[error("Invalid optimizer state: {0}")]
133    InvalidState(String),
134
135    /// Dimension mismatch
136    #[error("Dimension mismatch: expected {expected:?}, got {actual:?}")]
137    DimensionMismatch {
138        expected: Vec<usize>,
139        actual: Vec<usize>,
140    },
141
142    /// Not initialized
143    #[error("GPU optimizer not initialized")]
144    NotInitialized,
145
146    /// CUDA not available
147    #[error("CUDA is not available on this system")]
148    CudaNotAvailable,
149}
150
151/// Trait for GPU-accelerated optimizers
152pub trait GpuOptimizer<A: Float, D: Dimension> {
153    /// Check if GPU acceleration is available
154    fn is_gpu_available(&self) -> bool;
155
156    /// Move optimizer state to GPU
157    fn move_to_gpu(&mut self) -> Result<(), GpuOptimError>;
158
159    /// Move optimizer state back to CPU
160    fn move_to_cpu(&mut self) -> Result<(), GpuOptimError>;
161
162    /// Deprecated alias for [`Self::move_to_gpu`].
163    ///
164    /// `to_gpu` on a `&mut self` method triggers
165    /// `clippy::wrong_self_convention` (`to_*` names are conventionally
166    /// reserved for cheap `&self` -> owned conversions); `move_to_gpu`
167    /// names what this actually does. This shim delegates to
168    /// [`Self::move_to_gpu`] and exists only so 0.3.1-era callers keep
169    /// compiling.
170    #[deprecated(since = "0.3.2", note = "renamed to `move_to_gpu`")]
171    fn to_gpu(&mut self) -> Result<(), GpuOptimError> {
172        self.move_to_gpu()
173    }
174
175    /// Deprecated alias for [`Self::move_to_cpu`].
176    ///
177    /// See [`Self::to_gpu`] for why this was renamed. This shim delegates
178    /// to [`Self::move_to_cpu`] and exists only so 0.3.1-era callers keep
179    /// compiling.
180    #[deprecated(since = "0.3.2", note = "renamed to `move_to_cpu`")]
181    fn to_cpu(&mut self) -> Result<(), GpuOptimError> {
182        self.move_to_cpu()
183    }
184
185    /// Perform optimization step on GPU
186    fn step_gpu(
187        &mut self,
188        params: &mut Array<A, D>,
189        gradients: &Array<A, D>,
190    ) -> Result<(), GpuOptimError>;
191}
optirs_gpu/lib.rs

optirs_gpu/
lib.rs
