scirs2_core/gpu/kernels/
mod.rs

1//! GPU kernel library for common scientific computing operations
2//!
3//! This module provides optimized GPU kernels for various operations used in
4//! scientific computing, with support for multiple GPU backends.
5
6use std::collections::HashMap;
7use std::fmt;
8
9pub mod blas;
10pub mod complex;
11pub mod elementwise;
12pub mod ml;
13pub mod reduction;
14pub mod transform;
15
16use crate::gpu::{GpuBackend, GpuError};
17
18/// Supported data types for GPU kernels
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub enum DataType {
21    /// 32-bit floating point (f32)
22    Float32,
23    /// 64-bit floating point (f64)
24    Float64,
25    /// 32-bit signed integer (i32)
26    Int32,
27    /// 32-bit unsigned integer (u32)
28    UInt32,
29    /// 16-bit floating point (f16)
30    Float16,
31    /// Brain floating point (bfloat16)
32    BFloat16,
33}
34
35impl fmt::Display for DataType {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        match self {
38            DataType::Float32 => write!(f, "f32"),
39            DataType::Float64 => write!(f, "f64"),
40            DataType::Int32 => write!(f, "i32"),
41            DataType::UInt32 => write!(f, "u32"),
42            DataType::Float16 => write!(f, "f16"),
43            DataType::BFloat16 => write!(f, "bf16"),
44        }
45    }
46}
47
48/// The type of operation performed by the kernel
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum OperationType {
51    /// Primarily compute-intensive operations
52    ComputeIntensive,
53    /// Primarily memory-intensive operations
54    MemoryIntensive,
55    /// Balanced between compute and memory
56    Balanced,
57}
58
59/// Metadata for kernel execution
60#[derive(Debug, Clone)]
61pub struct KernelMetadata {
62    /// Recommended workgroup size
63    pub workgroup_size: [u32; 3],
64    /// Local memory usage in bytes
65    pub local_memory_usage: usize,
66    /// Whether the kernel supports tensor cores (NVIDIA) or similar
67    pub supports_tensor_cores: bool,
68    /// Operation type (compute intensive, memory intensive, balanced)
69    pub operationtype: OperationType,
70    /// Additional backend-specific metadata
71    pub backend_metadata: HashMap<String, String>,
72}
73
74impl Default for KernelMetadata {
75    fn default() -> Self {
76        Self {
77            workgroup_size: [16, 16, 1],
78            local_memory_usage: 0,
79            supports_tensor_cores: false,
80            operationtype: OperationType::Balanced,
81            backend_metadata: HashMap::new(),
82        }
83    }
84}
85
86/// Parameters for kernel specialization
87#[derive(Debug, Clone)]
88pub struct KernelParams {
89    /// Numeric type (f32, f64, etc.)
90    pub datatype: DataType,
91    /// Input dimensions
92    pub input_dims: Vec<usize>,
93    /// Output dimensions
94    pub output_dims: Vec<usize>,
95    /// Additional numeric parameters
96    pub numeric_params: HashMap<String, f64>,
97    /// Additional string parameters
98    pub string_params: HashMap<String, String>,
99}
100
101impl KernelParams {
102    /// Create new kernel parameters
103    pub fn new(datatype: DataType) -> Self {
104        Self {
105            datatype,
106            input_dims: Vec::new(),
107            output_dims: Vec::new(),
108            numeric_params: HashMap::new(),
109            string_params: HashMap::new(),
110        }
111    }
112
113    /// Set input dimensions
114    pub fn with_input_dims(mut self, dims: Vec<usize>) -> Self {
115        self.input_dims = dims;
116        self
117    }
118
119    /// Set output dimensions
120    pub fn with_output_dims(mut self, dims: Vec<usize>) -> Self {
121        self.output_dims = dims;
122        self
123    }
124
125    /// Add a numeric parameter
126    pub fn with_numeric_param(mut self, name: &str, value: f64) -> Self {
127        self.numeric_params.insert(name.to_string(), value);
128        self
129    }
130
131    /// Add a string parameter
132    pub fn with_string_param(mut self, name: &str, value: &str) -> Self {
133        self.string_params
134            .insert(name.to_string(), value.to_string());
135        self
136    }
137}
138
139/// GPU Kernel interface
140pub trait GpuKernel: Send + Sync {
141    /// The name of the kernel
142    fn name(&self) -> &str;
143
144    /// Get kernel source for the specified backend
145    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError>;
146
147    /// Get kernel metadata (workgroup size, memory requirements, etc.)
148    fn metadata(&self) -> KernelMetadata;
149
150    /// Can this kernel be specialized for the given parameters?
151    fn can_specialize(&self, params: &KernelParams) -> bool;
152
153    /// Create a specialized version of this kernel for the given parameters
154    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError>;
155}
156
157/// Base kernel implementation that can be used by specialized kernels
158pub struct BaseKernel {
159    name: String,
160    cuda_source: String,
161    rocm_source: String,
162    wgpu_source: String,
163    metal_source: String,
164    opencl_source: String,
165    metadata: KernelMetadata,
166}
167
168impl BaseKernel {
169    /// Create a new base kernel
170    pub fn new(
171        name: &str,
172        cuda_source: &str,
173        rocm_source: &str,
174        wgpu_source: &str,
175        metal_source: &str,
176        opencl_source: &str,
177        metadata: KernelMetadata,
178    ) -> Self {
179        Self {
180            name: name.to_string(),
181            cuda_source: cuda_source.to_string(),
182            rocm_source: rocm_source.to_string(),
183            wgpu_source: wgpu_source.to_string(),
184            metal_source: metal_source.to_string(),
185            opencl_source: opencl_source.to_string(),
186            metadata,
187        }
188    }
189}
190
191impl GpuKernel for BaseKernel {
192    fn name(&self) -> &str {
193        &self.name
194    }
195
196    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
197        match backend {
198            GpuBackend::Cuda => Ok(self.cuda_source.clone()),
199            GpuBackend::Rocm => Ok(self.rocm_source.clone()),
200            GpuBackend::Wgpu => Ok(self.wgpu_source.clone()),
201            GpuBackend::Metal => Ok(self.metal_source.clone()),
202            GpuBackend::OpenCL => Ok(self.opencl_source.clone()),
203            _ => Err(GpuError::UnsupportedBackend(backend)),
204        }
205    }
206
207    fn metadata(&self) -> KernelMetadata {
208        self.metadata.clone()
209    }
210
211    fn can_specialize(&self, params: &KernelParams) -> bool {
212        false // Base implementation doesn't support specialization
213    }
214
215    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
216        Err(GpuError::SpecializationNotSupported)
217    }
218}
219
220/// Registry of available GPU kernels
221pub struct KernelRegistry {
222    kernels: HashMap<String, Box<dyn GpuKernel>>,
223}
224
225impl KernelRegistry {
226    /// Create a new kernel registry
227    pub fn new() -> Self {
228        Self {
229            kernels: HashMap::new(),
230        }
231    }
232
233    /// Create a registry with all default kernels
234    pub fn with_default_kernels() -> Self {
235        let mut registry = Self::new();
236
237        // Register BLAS kernels
238        registry.register(Box::new(blas::gemm::GemmKernel::new()));
239        registry.register(Box::new(blas::axpy::AxpyKernel::new()));
240        registry.register(Box::new(blas::gemv::GemvKernel::new()));
241
242        // Register elementwise kernels
243        registry.register(Box::new(elementwise::ElementwiseAddKernel::new()));
244        registry.register(Box::new(elementwise::ElementwiseSubKernel::new()));
245        registry.register(Box::new(elementwise::ElementwiseMulKernel::new()));
246        registry.register(Box::new(elementwise::ElementwiseDivKernel::new()));
247        registry.register(Box::new(elementwise::ElementwisePowKernel::new()));
248        registry.register(Box::new(elementwise::ElementwiseSqrtKernel::new()));
249        registry.register(Box::new(elementwise::ElementwiseExpKernel::new()));
250        registry.register(Box::new(elementwise::ElementwiseLogKernel::new()));
251
252        // Register optimization kernels
253        registry.register(Box::new(create_adam_optimizer_kernel()));
254        registry.register(Box::new(create_sgd_optimizer_kernel()));
255        registry.register(Box::new(create_rmsprop_optimizer_kernel()));
256        registry.register(Box::new(create_adagrad_optimizer_kernel()));
257        registry.register(Box::new(create_lamb_optimizer_kernel()));
258
259        // Register utility kernels
260        registry.register(Box::new(create_memcpy_kernel()));
261        registry.register(Box::new(create_fill_kernel()));
262        registry.register(Box::new(create_reduce_sum_kernel()));
263        registry.register(Box::new(create_reduce_max_kernel()));
264
265        // Register transform kernels
266        registry.register(Box::new(transform::fft::FftKernel::new()));
267        registry.register(Box::new(transform::convolution::Conv1dKernel::new()));
268        registry.register(Box::new(transform::convolution::Conv2dKernel::new()));
269
270        // Register reduction kernels
271        registry.register(Box::new(reduction::sum::SumKernel::new()));
272        registry.register(Box::new(reduction::norm::NormKernel::new()));
273        registry.register(Box::new(reduction::min_max::MinKernel::new()));
274        registry.register(Box::new(reduction::min_max::MaxKernel::new()));
275        registry.register(Box::new(reduction::mean::MeanKernel::new()));
276        registry.register(Box::new(reduction::std_dev::StdDevKernel::new()));
277
278        // Register ML kernels
279        registry.register(Box::new(ml::activation::ReluKernel::new()));
280        registry.register(Box::new(ml::activation::SigmoidKernel::new()));
281        registry.register(Box::new(ml::activation::TanhKernel::new()));
282        registry.register(Box::new(ml::softmax::SoftmaxKernel::new()));
283        registry.register(Box::new(ml::pooling::MaxPoolKernel::new()));
284        registry.register(Box::new(ml::pooling::AvgPoolKernel::new()));
285
286        // Register complex number kernels
287        registry.register(Box::new(complex::ComplexMultiplyKernel::new()));
288        registry.register(Box::new(complex::ComplexConjugateKernel::new()));
289        registry.register(Box::new(complex::ComplexMatMulKernel::new()));
290
291        // Register RK4 integration kernels for advanced mode
292        registry.register(Box::new(create_rk4_stage1_kernel()));
293        registry.register(Box::new(create_rk4_stage2_kernel()));
294        registry.register(Box::new(create_rk4_stage3_kernel()));
295        registry.register(Box::new(create_rk4_stage4_kernel()));
296        registry.register(Box::new(create_rk4_combine_kernel()));
297        registry.register(Box::new(createerror_estimate_kernel()));
298
299        registry
300    }
301
302    /// Register a kernel
303    pub fn register(&mut self, kernel: Box<dyn GpuKernel>) {
304        self.kernels.insert(kernel.name().to_string(), kernel);
305    }
306
307    /// Get a kernel by name
308    pub fn get(&self, name: &str) -> Option<&dyn GpuKernel> {
309        self.kernels.get(name).map(|k| k.as_ref())
310    }
311
312    /// Get a specialized kernel
313    pub fn get_specialized(
314        &self,
315        name: &str,
316        params: &KernelParams,
317    ) -> Result<Box<dyn GpuKernel>, GpuError> {
318        let kernel = self
319            .get(name)
320            .ok_or_else(|| GpuError::KernelNotFound(name.to_string()))?;
321
322        if kernel.can_specialize(params) {
323            kernel.specialize(params)
324        } else {
325            Err(GpuError::SpecializationNotSupported)
326        }
327    }
328}
329
330impl Default for KernelRegistry {
331    fn default() -> Self {
332        Self::with_default_kernels()
333    }
334}
335
336// ─── WGSL shader sources for all kernels ─────────────────────────────────────
337
338/// WGSL source for the Adam optimizer kernel (workgroup 256).
339///
340/// Buffers (all at group 0):
341///   0 → params (read_write), 1 → grads (read), 2 → m (read_write),
342///   3 → v (read_write), 4 → uniforms (uniform)
343const ADAM_WGSL: &str = r#"
344@group(0) @binding(0) var<storage, read_write> params: array<f32>;
345@group(0) @binding(1) var<storage, read> grads: array<f32>;
346@group(0) @binding(2) var<storage, read_write> m: array<f32>;
347@group(0) @binding(3) var<storage, read_write> v: array<f32>;
348
349struct AdamUniforms {
350    lr: f32,
351    beta1: f32,
352    beta2: f32,
353    eps: f32,
354    weight_decay: f32,
355    bias_correction1: f32,
356    bias_correction2: f32,
357    n: u32,
358};
359
360@group(0) @binding(4) var<uniform> uniforms: AdamUniforms;
361
362@compute @workgroup_size(256)
363fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
364    let idx = global_id.x;
365    if idx >= uniforms.n { return; }
366
367    var grad = grads[idx];
368    if uniforms.weight_decay > 0.0 {
369        grad += uniforms.weight_decay * params[idx];
370    }
371
372    // Update biased first moment estimate
373    m[idx] = uniforms.beta1 * m[idx] + (1.0 - uniforms.beta1) * grad;
374
375    // Update biased second raw moment estimate
376    v[idx] = uniforms.beta2 * v[idx] + (1.0 - uniforms.beta2) * grad * grad;
377
378    // Bias-corrected moment estimates
379    let m_hat = m[idx] / uniforms.bias_correction1;
380    let v_hat = v[idx] / uniforms.bias_correction2;
381
382    // Parameter update
383    params[idx] -= uniforms.lr * m_hat / (sqrt(v_hat) + uniforms.eps);
384}
385"#;
386
387/// WGSL source for the SGD optimizer kernel (workgroup 256, with momentum).
388///
389/// Buffers: 0 → params (rw), 1 → grads (r), 2 → momentum_buf (rw)
390/// Uniforms: lr, momentum_factor, n
391const SGD_WGSL: &str = r#"
392@group(0) @binding(0) var<storage, read_write> params: array<f32>;
393@group(0) @binding(1) var<storage, read> grads: array<f32>;
394@group(0) @binding(2) var<storage, read_write> momentum_buf: array<f32>;
395
396struct SgdUniforms {
397    lr: f32,
398    momentum_factor: f32,
399    n: u32,
400};
401
402@group(0) @binding(3) var<uniform> uniforms: SgdUniforms;
403
404@compute @workgroup_size(256)
405fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
406    let idx = global_id.x;
407    if idx >= uniforms.n { return; }
408
409    let grad = grads[idx];
410    if uniforms.momentum_factor > 0.0 {
411        // SGD with momentum: buf = momentum * buf + grad; param -= lr * buf
412        let buf = uniforms.momentum_factor * momentum_buf[idx] + grad;
413        momentum_buf[idx] = buf;
414        params[idx] -= uniforms.lr * buf;
415    } else {
416        params[idx] -= uniforms.lr * grad;
417    }
418}
419"#;
420
421/// WGSL source for the RMSprop optimizer kernel (workgroup 256).
422///
423/// Buffers: 0 → params (rw), 1 → grads (r), 2 → cache (rw)
424/// Uniforms: lr, decay (alpha), epsilon, n
425const RMSPROP_WGSL: &str = r#"
426@group(0) @binding(0) var<storage, read_write> params: array<f32>;
427@group(0) @binding(1) var<storage, read> grads: array<f32>;
428@group(0) @binding(2) var<storage, read_write> cache: array<f32>;
429
430struct RmspropUniforms {
431    lr: f32,
432    decay: f32,
433    epsilon: f32,
434    n: u32,
435};
436
437@group(0) @binding(3) var<uniform> uniforms: RmspropUniforms;
438
439@compute @workgroup_size(256)
440fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
441    let idx = global_id.x;
442    if idx >= uniforms.n { return; }
443
444    let grad = grads[idx];
445    // cache = decay * cache + (1 - decay) * grad^2
446    let new_cache = uniforms.decay * cache[idx] + (1.0 - uniforms.decay) * grad * grad;
447    cache[idx] = new_cache;
448    // params -= lr * grad / (sqrt(cache) + epsilon)
449    params[idx] -= uniforms.lr * grad / (sqrt(new_cache) + uniforms.epsilon);
450}
451"#;
452
453/// WGSL source for the Adagrad optimizer kernel (workgroup 256).
454///
455/// Buffers: 0 → params (rw), 1 → grads (r), 2 → cache (rw, accumulated sq grads)
456/// Uniforms: lr, epsilon, n
457const ADAGRAD_WGSL: &str = r#"
458@group(0) @binding(0) var<storage, read_write> params: array<f32>;
459@group(0) @binding(1) var<storage, read> grads: array<f32>;
460@group(0) @binding(2) var<storage, read_write> cache: array<f32>;
461
462struct AdagradUniforms {
463    lr: f32,
464    epsilon: f32,
465    n: u32,
466};
467
468@group(0) @binding(3) var<uniform> uniforms: AdagradUniforms;
469
470@compute @workgroup_size(256)
471fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
472    let idx = global_id.x;
473    if idx >= uniforms.n { return; }
474
475    let grad = grads[idx];
476    // Accumulate squared gradient
477    let new_cache = cache[idx] + grad * grad;
478    cache[idx] = new_cache;
479    // Adaptive update
480    params[idx] -= uniforms.lr * grad / (sqrt(new_cache) + uniforms.epsilon);
481}
482"#;
483
484/// WGSL source for the LAMB optimizer kernel (workgroup 256, uniform-norm variant).
485///
486/// The caller pre-computes param_norm and grad_norm (L2 norms) and passes them
487/// as uniforms so a single pass can perform the layer-wise ratio scaling.
488///
489/// Buffers: 0 → params (rw), 1 → grads (r)
490/// Uniforms: lr, weight_decay, param_norm, grad_norm, n
491const LAMB_WGSL: &str = r#"
492@group(0) @binding(0) var<storage, read_write> params: array<f32>;
493@group(0) @binding(1) var<storage, read> grads: array<f32>;
494
495struct LambUniforms {
496    lr: f32,
497    weight_decay: f32,
498    param_norm: f32,
499    grad_norm: f32,
500    n: u32,
501};
502
503@group(0) @binding(2) var<uniform> uniforms: LambUniforms;
504
505@compute @workgroup_size(256)
506fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
507    let idx = global_id.x;
508    if idx >= uniforms.n { return; }
509
510    // update = grad + weight_decay * param
511    let update = grads[idx] + uniforms.weight_decay * params[idx];
512
513    // Layer-wise adaptive ratio: trust_ratio = param_norm / (grad_norm + eps)
514    // Guard against zero norms (use 1.0 as neutral ratio)
515    let eps = 1e-6;
516    let denom = uniforms.grad_norm + eps;
517    let trust_ratio = select(1.0, uniforms.param_norm / denom, uniforms.param_norm > 0.0 && uniforms.grad_norm > 0.0);
518
519    params[idx] -= uniforms.lr * trust_ratio * update;
520}
521"#;
522
523/// WGSL source for the memcpy kernel (workgroup 256).
524///
525/// Buffers: 0 → src (read), 1 → dst (read_write)
526/// Uniforms: n
527const MEMCPY_WGSL: &str = r#"
528@group(0) @binding(0) var<storage, read> src: array<f32>;
529@group(0) @binding(1) var<storage, read_write> dst: array<f32>;
530
531struct MemcpyUniforms {
532    n: u32,
533};
534
535@group(0) @binding(2) var<uniform> uniforms: MemcpyUniforms;
536
537@compute @workgroup_size(256)
538fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
539    let idx = global_id.x;
540    if idx >= uniforms.n { return; }
541    dst[idx] = src[idx];
542}
543"#;
544
545/// WGSL source for the fill kernel (workgroup 256).
546///
547/// Buffers: 0 → dst (read_write)
548/// Uniforms: value, n
549const FILL_WGSL: &str = r#"
550@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
551
552struct FillUniforms {
553    value: f32,
554    n: u32,
555};
556
557@group(0) @binding(1) var<uniform> uniforms: FillUniforms;
558
559@compute @workgroup_size(256)
560fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
561    let idx = global_id.x;
562    if idx >= uniforms.n { return; }
563    dst[idx] = uniforms.value;
564}
565"#;
566
567/// WGSL source for the reduce_sum kernel (single-pass workgroup reduction, workgroup 256).
568///
569/// Each workgroup reduces its slice into one partial sum written to `output[workgroup_id]`.
570/// The host must dispatch `ceil(n / 256)` workgroups and sum `output` on the CPU or with a
571/// second dispatch.
572///
573/// Buffers: 0 → input (read), 1 → output (read_write)
574/// Uniforms: n
575const REDUCE_SUM_WGSL: &str = r#"
576@group(0) @binding(0) var<storage, read> input: array<f32>;
577@group(0) @binding(1) var<storage, read_write> output: array<f32>;
578
579struct ReduceUniforms {
580    n: u32,
581};
582
583@group(0) @binding(2) var<uniform> uniforms: ReduceUniforms;
584
585var<workgroup> scratch: array<f32, 256>;
586
587@compute @workgroup_size(256)
588fn main(
589    @builtin(global_invocation_id) global_id: vec3<u32>,
590    @builtin(local_invocation_id)  local_id:  vec3<u32>,
591    @builtin(workgroup_id)         wg_id:     vec3<u32>,
592) {
593    let gidx = global_id.x;
594    let lidx = local_id.x;
595
596    // Load with bounds guard
597    if gidx < uniforms.n {
598        scratch[lidx] = input[gidx];
599    } else {
600        scratch[lidx] = 0.0;
601    }
602    workgroupBarrier();
603
604    // Tree reduction within the workgroup
605    var stride = 128u;
606    loop {
607        if stride == 0u { break; }
608        if lidx < stride {
609            scratch[lidx] += scratch[lidx + stride];
610        }
611        workgroupBarrier();
612        if stride == 1u { break; }
613        stride = stride >> 1u;
614    }
615
616    // Thread 0 writes the partial sum for this workgroup
617    if lidx == 0u {
618        output[wg_id.x] = scratch[0];
619    }
620}
621"#;
622
623/// WGSL source for the reduce_max kernel (single-pass workgroup reduction, workgroup 256).
624///
625/// Uses the same two-pass convention as reduce_sum; each workgroup writes a partial maximum.
626///
627/// Buffers: 0 → input (read), 1 → output (read_write)
628/// Uniforms: n
629const REDUCE_MAX_WGSL: &str = r#"
630@group(0) @binding(0) var<storage, read> input: array<f32>;
631@group(0) @binding(1) var<storage, read_write> output: array<f32>;
632
633struct ReduceUniforms {
634    n: u32,
635};
636
637@group(0) @binding(2) var<uniform> uniforms: ReduceUniforms;
638
639var<workgroup> scratch: array<f32, 256>;
640
641@compute @workgroup_size(256)
642fn main(
643    @builtin(global_invocation_id) global_id: vec3<u32>,
644    @builtin(local_invocation_id)  local_id:  vec3<u32>,
645    @builtin(workgroup_id)         wg_id:     vec3<u32>,
646) {
647    let gidx = global_id.x;
648    let lidx = local_id.x;
649
650    // Load with bounds guard; use -f32::MAX as neutral element for max
651    if gidx < uniforms.n {
652        scratch[lidx] = input[gidx];
653    } else {
654        scratch[lidx] = -3.402823e+38; // -FLT_MAX
655    }
656    workgroupBarrier();
657
658    // Tree reduction within the workgroup
659    var stride = 128u;
660    loop {
661        if stride == 0u { break; }
662        if lidx < stride {
663            scratch[lidx] = max(scratch[lidx], scratch[lidx + stride]);
664        }
665        workgroupBarrier();
666        if stride == 1u { break; }
667        stride = stride >> 1u;
668    }
669
670    if lidx == 0u {
671        output[wg_id.x] = scratch[0];
672    }
673}
674"#;
675
676/// WGSL for RK4 stage 1: k1[i] = h * f(t, y[i])  where f(t, y) = -y  (exponential decay placeholder).
677///
678/// Buffers: 0 → y (read), 1 → k1 (read_write)
679/// Uniforms: t, h, n
680const RK4_STAGE1_WGSL: &str = r#"
681@group(0) @binding(0) var<storage, read> y: array<f32>;
682@group(0) @binding(1) var<storage, read_write> k1: array<f32>;
683
684struct Rk4Uniforms {
685    t: f32,
686    h: f32,
687    n: u32,
688};
689
690@group(0) @binding(2) var<uniform> uniforms: Rk4Uniforms;
691
692@compute @workgroup_size(256)
693fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
694    let idx = global_id.x;
695    if idx >= uniforms.n { return; }
696    // Placeholder ODE: dy/dt = -y  (exponential decay)
697    let dydt = -y[idx];
698    k1[idx] = uniforms.h * dydt;
699}
700"#;
701
702/// WGSL for RK4 stage 2: k2[i] = h * f(t + h/2, y[i] + k1[i]/2).
703///
704/// Buffers: 0 → y (read), 1 → k1 (read), 2 → k2 (read_write)
705/// Uniforms: t, h, n
706const RK4_STAGE2_WGSL: &str = r#"
707@group(0) @binding(0) var<storage, read> y: array<f32>;
708@group(0) @binding(1) var<storage, read> k1: array<f32>;
709@group(0) @binding(2) var<storage, read_write> k2: array<f32>;
710
711struct Rk4Uniforms {
712    t: f32,
713    h: f32,
714    n: u32,
715};
716
717@group(0) @binding(3) var<uniform> uniforms: Rk4Uniforms;
718
719@compute @workgroup_size(256)
720fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
721    let idx = global_id.x;
722    if idx >= uniforms.n { return; }
723    let y_mid = y[idx] + 0.5 * k1[idx];
724    // Placeholder ODE: dy/dt = -y
725    let dydt = -y_mid;
726    k2[idx] = uniforms.h * dydt;
727}
728"#;
729
730/// WGSL for RK4 stage 3: k3[i] = h * f(t + h/2, y[i] + k2[i]/2).
731///
732/// Buffers: 0 → y (read), 1 → k2 (read), 2 → k3 (read_write)
733/// Uniforms: t, h, n
734const RK4_STAGE3_WGSL: &str = r#"
735@group(0) @binding(0) var<storage, read> y: array<f32>;
736@group(0) @binding(1) var<storage, read> k2: array<f32>;
737@group(0) @binding(2) var<storage, read_write> k3: array<f32>;
738
739struct Rk4Uniforms {
740    t: f32,
741    h: f32,
742    n: u32,
743};
744
745@group(0) @binding(3) var<uniform> uniforms: Rk4Uniforms;
746
747@compute @workgroup_size(256)
748fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
749    let idx = global_id.x;
750    if idx >= uniforms.n { return; }
751    let y_mid = y[idx] + 0.5 * k2[idx];
752    // Placeholder ODE: dy/dt = -y
753    let dydt = -y_mid;
754    k3[idx] = uniforms.h * dydt;
755}
756"#;
757
758/// WGSL for RK4 stage 4: k4[i] = h * f(t + h, y[i] + k3[i]).
759///
760/// Buffers: 0 → y (read), 1 → k3 (read), 2 → k4 (read_write)
761/// Uniforms: t, h, n
762const RK4_STAGE4_WGSL: &str = r#"
763@group(0) @binding(0) var<storage, read> y: array<f32>;
764@group(0) @binding(1) var<storage, read> k3: array<f32>;
765@group(0) @binding(2) var<storage, read_write> k4: array<f32>;
766
767struct Rk4Uniforms {
768    t: f32,
769    h: f32,
770    n: u32,
771};
772
773@group(0) @binding(3) var<uniform> uniforms: Rk4Uniforms;
774
775@compute @workgroup_size(256)
776fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
777    let idx = global_id.x;
778    if idx >= uniforms.n { return; }
779    let y_next = y[idx] + k3[idx];
780    // Placeholder ODE: dy/dt = -y
781    let dydt = -y_next;
782    k4[idx] = uniforms.h * dydt;
783}
784"#;
785
786/// WGSL for RK4 final combination: y_new[i] = y[i] + (k1 + 2*k2 + 2*k3 + k4) / 6.
787///
788/// Buffers: 0 → y (read), 1 → k1 (read), 2 → k2 (read), 3 → k3 (read), 4 → k4 (read),
789///          5 → y_new (read_write)
790/// Uniforms: n
791const RK4_COMBINE_WGSL: &str = r#"
792@group(0) @binding(0) var<storage, read> y: array<f32>;
793@group(0) @binding(1) var<storage, read> k1: array<f32>;
794@group(0) @binding(2) var<storage, read> k2: array<f32>;
795@group(0) @binding(3) var<storage, read> k3: array<f32>;
796@group(0) @binding(4) var<storage, read> k4: array<f32>;
797@group(0) @binding(5) var<storage, read_write> y_new: array<f32>;
798
799struct Rk4CombineUniforms {
800    n: u32,
801};
802
803@group(0) @binding(6) var<uniform> uniforms: Rk4CombineUniforms;
804
805@compute @workgroup_size(256)
806fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
807    let idx = global_id.x;
808    if idx >= uniforms.n { return; }
809    let weighted = k1[idx] + 2.0 * k2[idx] + 2.0 * k3[idx] + k4[idx];
810    y_new[idx] = y[idx] + weighted * (1.0 / 6.0);
811}
812"#;
813
814/// WGSL for the error estimate kernel: err[i] = |y1[i] - y2[i]| / max(scale, eps).
815///
816/// scale = atol + rtol * max(|y1|, |y2|) — matches CUDA error_estimate.cu semantics.
817///
818/// Buffers: 0 → y1 (read), 1 → y2 (read), 2 → err (read_write)
819/// Uniforms: rtol, atol, n
820const ERROR_ESTIMATE_WGSL: &str = r#"
821@group(0) @binding(0) var<storage, read> y1: array<f32>;
822@group(0) @binding(1) var<storage, read> y2: array<f32>;
823@group(0) @binding(2) var<storage, read_write> err: array<f32>;
824
825struct ErrorUniforms {
826    rtol: f32,
827    atol: f32,
828    n: u32,
829};
830
831@group(0) @binding(3) var<uniform> uniforms: ErrorUniforms;
832
833@compute @workgroup_size(256)
834fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
835    let idx = global_id.x;
836    if idx >= uniforms.n { return; }
837
838    let v1 = y1[idx];
839    let v2 = y2[idx];
840    let abs_err = abs(v1 - v2);
841    let y_scale = max(abs(v1), abs(v2));
842    let scale = uniforms.atol + uniforms.rtol * y_scale;
843    err[idx] = abs_err / max(scale, 1e-7);
844}
845"#;
846
847// ─── Kernel factory functions ─────────────────────────────────────────────────
848
849/// Create RK4 Stage 1 kernel for advanced mode GPU acceleration
850fn create_rk4_stage1_kernel() -> BaseKernel {
851    let cuda_source = include_str!("rk4_stage1.cu");
852    let metadata = KernelMetadata {
853        workgroup_size: [256, 1, 1],
854        local_memory_usage: 0,
855        supports_tensor_cores: false,
856        operationtype: OperationType::ComputeIntensive,
857        backend_metadata: HashMap::new(),
858    };
859
860    BaseKernel::new(
861        "rk4_stage1",
862        cuda_source,
863        cuda_source, // Use CUDA source for ROCm (HIP compatible)
864        RK4_STAGE1_WGSL,
865        "",          // Metal source not implemented yet
866        cuda_source, // Use CUDA source for OpenCL (with minor modifications)
867        metadata,
868    )
869}
870
871/// Create RK4 Stage 2 kernel for advanced mode GPU acceleration
872fn create_rk4_stage2_kernel() -> BaseKernel {
873    let cuda_source = include_str!("rk4_stage2.cu");
874    let metadata = KernelMetadata {
875        workgroup_size: [256, 1, 1],
876        local_memory_usage: 0,
877        supports_tensor_cores: false,
878        operationtype: OperationType::ComputeIntensive,
879        backend_metadata: HashMap::new(),
880    };
881
882    BaseKernel::new(
883        "rk4_stage2",
884        cuda_source,
885        cuda_source,
886        RK4_STAGE2_WGSL,
887        "",
888        cuda_source,
889        metadata,
890    )
891}
892
893/// Create RK4 Stage 3 kernel for advanced mode GPU acceleration
894fn create_rk4_stage3_kernel() -> BaseKernel {
895    let cuda_source = include_str!("rk4_stage3.cu");
896    let metadata = KernelMetadata {
897        workgroup_size: [256, 1, 1],
898        local_memory_usage: 0,
899        supports_tensor_cores: false,
900        operationtype: OperationType::ComputeIntensive,
901        backend_metadata: HashMap::new(),
902    };
903
904    BaseKernel::new(
905        "rk4_stage3",
906        cuda_source,
907        cuda_source,
908        RK4_STAGE3_WGSL,
909        "",
910        cuda_source,
911        metadata,
912    )
913}
914
915/// Create RK4 Stage 4 kernel for advanced mode GPU acceleration
916fn create_rk4_stage4_kernel() -> BaseKernel {
917    let cuda_source = include_str!("rk4_stage4.cu");
918    let metadata = KernelMetadata {
919        workgroup_size: [256, 1, 1],
920        local_memory_usage: 0,
921        supports_tensor_cores: false,
922        operationtype: OperationType::ComputeIntensive,
923        backend_metadata: HashMap::new(),
924    };
925
926    BaseKernel::new(
927        "rk4_stage4",
928        cuda_source,
929        cuda_source,
930        RK4_STAGE4_WGSL,
931        "",
932        cuda_source,
933        metadata,
934    )
935}
936
937/// Create RK4 Combination kernel for advanced mode GPU acceleration
938fn create_rk4_combine_kernel() -> BaseKernel {
939    let cuda_source = include_str!("rk4_combine.cu");
940    let metadata = KernelMetadata {
941        workgroup_size: [256, 1, 1],
942        local_memory_usage: 0,
943        supports_tensor_cores: false,
944        operationtype: OperationType::MemoryIntensive,
945        backend_metadata: HashMap::new(),
946    };
947
948    BaseKernel::new(
949        "rk4_combine",
950        cuda_source,
951        cuda_source,
952        RK4_COMBINE_WGSL,
953        "",
954        cuda_source,
955        metadata,
956    )
957}
958
959/// Create Error Estimation kernel for adaptive step size control
960fn createerror_estimate_kernel() -> BaseKernel {
961    let cuda_source = include_str!("error_estimate.cu");
962    let metadata = KernelMetadata {
963        workgroup_size: [256, 1, 1],
964        local_memory_usage: 1024, // Shared memory for reduction
965        supports_tensor_cores: false,
966        operationtype: OperationType::ComputeIntensive,
967        backend_metadata: HashMap::new(),
968    };
969
970    BaseKernel::new(
971        "error_estimate",
972        cuda_source,
973        cuda_source,
974        ERROR_ESTIMATE_WGSL,
975        "",
976        cuda_source,
977        metadata,
978    )
979}
980
981/// Create Adam optimizer kernel for GPU acceleration
982fn create_adam_optimizer_kernel() -> BaseKernel {
983    let cuda_source = include_str!("adam_optimizer.cu");
984    let metadata = KernelMetadata {
985        workgroup_size: [256, 1, 1],
986        local_memory_usage: 0,
987        supports_tensor_cores: false,
988        operationtype: OperationType::ComputeIntensive,
989        backend_metadata: HashMap::new(),
990    };
991
992    BaseKernel::new(
993        "adam_optimizer",
994        cuda_source,
995        cuda_source,
996        ADAM_WGSL,
997        "",
998        cuda_source,
999        metadata,
1000    )
1001}
1002
1003/// Create SGD optimizer kernel for GPU acceleration
1004fn create_sgd_optimizer_kernel() -> BaseKernel {
1005    let cuda_source = include_str!("sgd_optimizer.cu");
1006    let metadata = KernelMetadata {
1007        workgroup_size: [256, 1, 1],
1008        local_memory_usage: 0,
1009        supports_tensor_cores: false,
1010        operationtype: OperationType::MemoryIntensive,
1011        backend_metadata: HashMap::new(),
1012    };
1013
1014    BaseKernel::new(
1015        "sgd_optimizer",
1016        cuda_source,
1017        cuda_source,
1018        SGD_WGSL,
1019        "",
1020        cuda_source,
1021        metadata,
1022    )
1023}
1024
1025/// Create RMSprop optimizer kernel for GPU acceleration
1026fn create_rmsprop_optimizer_kernel() -> BaseKernel {
1027    let cuda_source = include_str!("rmsprop_optimizer.cu");
1028    let metadata = KernelMetadata {
1029        workgroup_size: [256, 1, 1],
1030        local_memory_usage: 0,
1031        supports_tensor_cores: false,
1032        operationtype: OperationType::ComputeIntensive,
1033        backend_metadata: HashMap::new(),
1034    };
1035
1036    BaseKernel::new(
1037        "rmsprop_optimizer",
1038        cuda_source,
1039        cuda_source,
1040        RMSPROP_WGSL,
1041        "",
1042        cuda_source,
1043        metadata,
1044    )
1045}
1046
1047/// Create Adagrad optimizer kernel for GPU acceleration
1048fn create_adagrad_optimizer_kernel() -> BaseKernel {
1049    let cuda_source = include_str!("adagrad_optimizer.cu");
1050    let metadata = KernelMetadata {
1051        workgroup_size: [256, 1, 1],
1052        local_memory_usage: 0,
1053        supports_tensor_cores: false,
1054        operationtype: OperationType::ComputeIntensive,
1055        backend_metadata: HashMap::new(),
1056    };
1057
1058    BaseKernel::new(
1059        "adagrad_optimizer",
1060        cuda_source,
1061        cuda_source,
1062        ADAGRAD_WGSL,
1063        "",
1064        cuda_source,
1065        metadata,
1066    )
1067}
1068
1069/// Create LAMB optimizer kernel for GPU acceleration
1070fn create_lamb_optimizer_kernel() -> BaseKernel {
1071    let cuda_source = include_str!("lamb_optimizer.cu");
1072    let metadata = KernelMetadata {
1073        workgroup_size: [256, 1, 1],
1074        local_memory_usage: 0,
1075        supports_tensor_cores: false,
1076        operationtype: OperationType::ComputeIntensive,
1077        backend_metadata: HashMap::new(),
1078    };
1079
1080    BaseKernel::new(
1081        "lamb_optimizer",
1082        cuda_source,
1083        cuda_source,
1084        LAMB_WGSL,
1085        "",
1086        cuda_source,
1087        metadata,
1088    )
1089}
1090
1091/// Create memory copy kernel for GPU acceleration
1092fn create_memcpy_kernel() -> BaseKernel {
1093    let cuda_source = include_str!("memcpy.cu");
1094    let metadata = KernelMetadata {
1095        workgroup_size: [256, 1, 1],
1096        local_memory_usage: 0,
1097        supports_tensor_cores: false,
1098        operationtype: OperationType::MemoryIntensive,
1099        backend_metadata: HashMap::new(),
1100    };
1101
1102    BaseKernel::new(
1103        "memcpy",
1104        cuda_source,
1105        cuda_source,
1106        MEMCPY_WGSL,
1107        "",
1108        cuda_source,
1109        metadata,
1110    )
1111}
1112
1113/// Create fill kernel for GPU acceleration
1114fn create_fill_kernel() -> BaseKernel {
1115    let cuda_source = include_str!("fill.cu");
1116    let metadata = KernelMetadata {
1117        workgroup_size: [256, 1, 1],
1118        local_memory_usage: 0,
1119        supports_tensor_cores: false,
1120        operationtype: OperationType::MemoryIntensive,
1121        backend_metadata: HashMap::new(),
1122    };
1123
1124    BaseKernel::new(
1125        "fill",
1126        cuda_source,
1127        cuda_source,
1128        FILL_WGSL,
1129        "",
1130        cuda_source,
1131        metadata,
1132    )
1133}
1134
1135/// Create reduce sum kernel for GPU acceleration
1136fn create_reduce_sum_kernel() -> BaseKernel {
1137    let cuda_source = include_str!("reduce_sum.cu");
1138    let metadata = KernelMetadata {
1139        workgroup_size: [256, 1, 1],
1140        local_memory_usage: 1024, // Shared memory for reduction
1141        supports_tensor_cores: false,
1142        operationtype: OperationType::ComputeIntensive,
1143        backend_metadata: HashMap::new(),
1144    };
1145
1146    BaseKernel::new(
1147        "reduce_sum",
1148        cuda_source,
1149        cuda_source,
1150        REDUCE_SUM_WGSL,
1151        "",
1152        cuda_source,
1153        metadata,
1154    )
1155}
1156
1157/// Create reduce max kernel for GPU acceleration
1158fn create_reduce_max_kernel() -> BaseKernel {
1159    let cuda_source = include_str!("reduce_max.cu");
1160    let metadata = KernelMetadata {
1161        workgroup_size: [256, 1, 1],
1162        local_memory_usage: 1024, // Shared memory for reduction
1163        supports_tensor_cores: false,
1164        operationtype: OperationType::ComputeIntensive,
1165        backend_metadata: HashMap::new(),
1166    };
1167
1168    BaseKernel::new(
1169        "reduce_max",
1170        cuda_source,
1171        cuda_source,
1172        REDUCE_MAX_WGSL,
1173        "",
1174        cuda_source,
1175        metadata,
1176    )
1177}
scirs2_core/gpu/kernels/mod.rs

scirs2_core/gpu/kernels/
mod.rs
