Crate optirs_gpu 

OptiRS GPU — GPU acceleration and GPU-aware optimizer tooling

Version: 0.3.2

optirs-gpu has two halves, and it is worth being precise about which is which:

1. A real GPU optimizer path. optimizers runs Adam, AdamW, SGD, RMSprop, Adagrad and LAMB as compute shaders through scirs2_core::gpu. Parameters and gradients are uploaded to device buffers, a compiled pipeline is dispatched, and the result is read back; the per-parameter optimizer state stays resident in device memory between steps. The kernels ship in both WGSL and MSL (shaders) so the same optimizer runs on whichever backend the machine can reach.
2. A CPU library of GPU-aware algorithms. occupancy, kernel_fusion, quantization, sparse_optimizer and memory::allocation / memory::management are pure-CPU models, planners and numerical routines that reason about GPU execution. They have no device dependency and are fully covered by unit tests.

Backend support matrix

Backend	Status
Metal (metal, automatic on macOS)	✅ real compute: MSL pipelines, buffers, dispatch, readback
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
OpenCL (opencl)	🚧 context creation only — no OpenCL C kernel sources are shipped
CUDA (cuda)	❌ not available — scirs2-core removed its CUDA backend in 0.6.x; the feature gates reporting code only
ROCm	❌ not available

optimizers::GpuOptimizerConfig defaults to probing optimizers::SUPPORTED_BACKENDS in order and using the first that opens.

Not implemented (and not faked)

• Cross-device collectives. multi_gpu can drive a real reduction kernel on a single device; anything that would require moving data between two physical GPUs returns GpuOptimError::UnsupportedOperation.
• Literal NVIDIA tensor cores / wmma. tensor_cores provides real CPU-side matrix-layout optimization, precision selection and AMP loss scaling; the device GEMM entry points (tensor_core_gemm, fused_adam_tensor_core, …) report an honest GpuOptimError::UnsupportedOperation because no backend this crate can reach exposes NVIDIA tensor cores.

Example

use optirs_gpu::optimizers::{AdamParams, GpuAdam};
use optirs_gpu::GpuOptimizer;
use scirs2_core::ndarray::Array1;

let mut optimizer = GpuAdam::new(AdamParams::default())?;
optimizer.move_to_gpu()?;

let mut params = Array1::from_elem(1_024, 1.0f32);
let grads = Array1::from_elem(1_024, 0.01f32);
optimizer.step_gpu(&mut params, &grads)?;

// Bring the moment estimates back to host memory when done.
optimizer.move_to_cpu()?;

Architecture

Every device access goes through SciRS2:

• GPU context: scirs2_core::gpu::GpuContext
• GPU memory: scirs2_core::gpu::GpuBuffer
• Kernel compilation: scirs2_core::gpu::GpuCompiler

wgpu, pollster, metal and friends are not direct dependencies of this crate; they arrive through the scirs2-core/<backend> features that this crate’s features forward.

Re-exports

pub use backends::GpuBackend;

pub use kernel_fusion::FusionGraph;

pub use kernel_fusion::FusionGroup;

pub use kernel_fusion::FusionOp;

pub use kernel_fusion::FusionPlan;

pub use kernel_fusion::FusionPlanner;

pub use kernel_fusion::OpKind;

pub use memory::MemoryPool;

pub use mixed_precision::f16_bits_to_f32;

pub use mixed_precision::f32_to_f16_bits;

pub use mixed_precision::DynamicLossScaler;

pub use mixed_precision::MixedPrecisionConfig;

pub use mixed_precision::OverflowStats;

pub use occupancy::calculate_occupancy;

pub use occupancy::optimal_block_size;

pub use occupancy::KernelResourceUsage;

pub use occupancy::OccupancyLimiter;

pub use occupancy::OccupancyResult;

pub use occupancy::SmResourceLimits;

pub use optimizers::AdagradParams;

pub use optimizers::AdamParams;

pub use optimizers::GpuAdagrad;

pub use optimizers::GpuAdam;

pub use optimizers::GpuAdamW;

pub use optimizers::GpuLamb;

pub use optimizers::GpuOptimizerConfig;

pub use optimizers::GpuRmsprop;

pub use optimizers::GpuSgd;

pub use optimizers::RmspropParams;

pub use optimizers::SgdParams;

pub use quantization::fake_quant_backward;

pub use quantization::fake_quant_fp8;

pub use quantization::fake_quant_int;

pub use quantization::fake_quant_int_per_channel;

pub use quantization::per_channel_params;

pub use quantization::Fp8Format;

pub use quantization::IntDtype;

pub use quantization::QatConfig;

pub use quantization::QatOptimizer;

pub use quantization::QuantParams;

pub use quantization::QuantScheme;

pub use quantization::QuantTarget;

pub use quantization::RoundingMode;

pub use sparse_optimizer::CooGradient;

pub use sparse_optimizer::CsrGradient;

pub use sparse_optimizer::LazyAdamMode;

pub use sparse_optimizer::SparseAdam;

pub use sparse_optimizer::SparseAdamConfig;

pub use sparse_optimizer::SparseAdamTable;

pub use sparse_optimizer::SparseSgd;

pub use sparse_optimizer::SparseSgdConfig;

pub use sparse_optimizer::SparseSgdTable;

Modules

backends

GPU backend identifiers and capability data.

kernel_fusion

Elementwise Kernel-Fusion Planner

memory

mixed_precision

Mixed-precision (AMP) building blocks: IEEE-754 binary16 conversion and dynamic loss scaling.

multi_gpu

occupancy

CUDA Occupancy Calculator (CPU-side, analytical).

optimizers

GPU-resident optimizer steps executed through scirs2_core::gpu.

quantization

Quantization-Aware Training (QAT) primitives

shaders

Compute-shader sources for the GPU optimizer steps, in WGSL and MSL.

sparse_optimizer

Sparse optimizer step (CSR/COO + lazy-state sparse SGD / Adam)

tensor_cores

Auto-generated module structure

utils

Small utilities shared by the GPU memory and kernel-launch paths.

Enums

GpuOptimError

Error type for GPU optimizer operations

Traits

GpuOptimizer

Trait for GPU-accelerated optimizers