tenflowers_core/
dispatch_init.rs

1/// Dispatch Registry Initialization Module
2///
3/// This module handles initialization of all operation registrations for the
4/// unified dispatch system. It should be called once at library initialization.
5///
6/// The initialization is performed lazily using lazy_static to ensure it happens
7/// exactly once before any dispatch calls are made.
8use crate::dispatch_registry_examples::initialize_dispatch_registrations;
9use crate::dispatch_registry_extended::initialize_extended_registrations;
10
11use lazy_static::lazy_static;
12
13lazy_static! {
14    /// Ensures dispatch registry is initialized before first use
15    pub static ref DISPATCH_INIT: () = {
16        initialize_dispatch_registrations();
17        initialize_extended_registrations();
18        #[cfg(feature = "gpu")]
19        initialize_gpu_operations();
20    };
21}
22
23/// Initialize GPU-specific operations
24#[cfg(feature = "gpu")]
25fn initialize_gpu_operations() {
26    // Reduction operations (sum, mean, ...) are registered here with a real CPU
27    // kernel plus a GPU kernel. The CPU kernel keeps CPU-tensor dispatch working
28    // (and correct) even when no usable GPU device exists.
29    register_gpu_reductions();
30}
31
32/// Register GPU reduction operations
33#[cfg(feature = "gpu")]
34fn register_gpu_reductions() {
35    use crate::dispatch_registry::{
36        BackendType, KernelImplementation, OperationDescriptor, F32_REGISTRY,
37    };
38    use crate::DType;
39
40    // Sum operation for f32
41    {
42        let desc = OperationDescriptor::new("sum", "reduction").with_dtypes(vec![DType::Float32]);
43
44        // Only register if not already registered (from examples)
45        if F32_REGISTRY.get_operation("sum").is_none() {
46            F32_REGISTRY.register_operation(desc).ok();
47        }
48
49        // Register the real CPU kernel so `sum` dispatches (and computes) on CPU
50        // tensors even when no usable GPU device exists. Without it a CPU tensor
51        // would fall back to the GPU kernel and fail with "device lost" on
52        // headless / sandboxed hosts.
53        F32_REGISTRY
54            .register_kernel(
55                "sum",
56                KernelImplementation::unary(BackendType::Cpu, sum_f32_cpu),
57            )
58            .ok();
59
60        // Register GPU kernel (used for GPU-resident tensors)
61        F32_REGISTRY
62            .register_kernel(
63                "sum",
64                KernelImplementation::unary(BackendType::Gpu, sum_f32_gpu),
65            )
66            .ok();
67    }
68
69    // Mean operation for f32
70    {
71        let desc = OperationDescriptor::new("mean", "reduction").with_dtypes(vec![DType::Float32]);
72
73        if F32_REGISTRY.get_operation("mean").is_none() {
74            F32_REGISTRY.register_operation(desc).ok();
75        }
76
77        // Real CPU kernel (see the `sum` note above).
78        F32_REGISTRY
79            .register_kernel(
80                "mean",
81                KernelImplementation::unary(BackendType::Cpu, mean_f32_cpu),
82            )
83            .ok();
84
85        // Register GPU kernel (used for GPU-resident tensors)
86        F32_REGISTRY
87            .register_kernel(
88                "mean",
89                KernelImplementation::unary(BackendType::Gpu, mean_f32_gpu),
90            )
91            .ok();
92    }
93}
94
95/// GPU implementation of sum for f32
96#[cfg(feature = "gpu")]
97fn sum_f32_gpu(x: &crate::Tensor<f32>) -> crate::Result<crate::Tensor<f32>> {
98    use crate::gpu::buffer::GpuBuffer;
99    use crate::gpu::ops::operation_types::ReductionOp;
100    use crate::gpu::ops::reduction_ops::execute_reduction_op;
101    use crate::Device;
102    use scirs2_core::ndarray::Array;
103
104    // Get tensor data as slice
105    let slice = x.data();
106
107    // Create GPU buffer from slice
108    let gpu_buffer = GpuBuffer::from_slice(slice, &Device::Gpu(0))?;
109
110    // Execute GPU reduction
111    let result_buffer = execute_reduction_op(&gpu_buffer, ReductionOp::Sum, None)?;
112
113    // Read back result from GPU
114    let result_data = result_buffer.to_cpu()?;
115
116    // Convert to tensor
117    let result = Array::from_elem(vec![], result_data[0]).into_dyn();
118    Ok(crate::Tensor::from_array(result))
119}
120
121/// GPU implementation of mean for f32
122#[cfg(feature = "gpu")]
123fn mean_f32_gpu(x: &crate::Tensor<f32>) -> crate::Result<crate::Tensor<f32>> {
124    use crate::gpu::buffer::GpuBuffer;
125    use crate::gpu::ops::operation_types::ReductionOp;
126    use crate::gpu::ops::reduction_ops::execute_reduction_op;
127    use crate::Device;
128    use scirs2_core::ndarray::Array;
129
130    // Get tensor data as slice
131    let slice = x.data();
132
133    // Create GPU buffer from slice
134    let gpu_buffer = GpuBuffer::from_slice(slice, &Device::Gpu(0))?;
135
136    // Execute GPU reduction
137    let result_buffer = execute_reduction_op(&gpu_buffer, ReductionOp::Mean, None)?;
138
139    // Read back result from GPU
140    let result_data = result_buffer.to_cpu()?;
141
142    // Convert to tensor
143    let result = Array::from_elem(vec![], result_data[0]).into_dyn();
144    Ok(crate::Tensor::from_array(result))
145}
146
147/// CPU fallback for sum
148fn sum_f32_cpu(x: &crate::Tensor<f32>) -> crate::Result<crate::Tensor<f32>> {
149    use scirs2_core::ndarray::Array;
150
151    let data = x.data();
152    let sum: f32 = data.iter().sum();
153    let result = Array::from_elem(vec![], sum).into_dyn();
154    Ok(crate::Tensor::from_array(result))
155}
156
157/// CPU fallback for mean
158fn mean_f32_cpu(x: &crate::Tensor<f32>) -> crate::Result<crate::Tensor<f32>> {
159    use scirs2_core::ndarray::Array;
160
161    let data = x.data();
162    let sum: f32 = data.iter().sum();
163    let mean = sum / data.len() as f32;
164    let result = Array::from_elem(vec![], mean).into_dyn();
165    Ok(crate::Tensor::from_array(result))
166}
167
168/// Ensure dispatch registry is initialized
169///
170/// This function can be called explicitly to ensure initialization has occurred.
171/// It's automatically called on first access to global registries, but can be
172/// called early for better control over initialization timing.
173pub fn ensure_initialized() {
174    *DISPATCH_INIT;
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::Tensor;
181    use scirs2_core::ndarray::array;
182
183    #[test]
184    fn test_initialization() {
185        ensure_initialized();
186
187        use crate::dispatch_registry::F32_REGISTRY;
188
189        // Verify operations are registered
190        assert!(F32_REGISTRY.get_operation("add").is_some());
191        assert!(F32_REGISTRY.get_operation("mul").is_some());
192        assert!(F32_REGISTRY.get_operation("div").is_some());
193    }
194
195    #[test]
196    fn test_sum_cpu_fallback() {
197        let input = Tensor::from_array(array![1.0f32, 2.0, 3.0, 4.0, 5.0].into_dyn());
198        let result = sum_f32_cpu(&input).expect("test: sum_f32_cpu should succeed");
199
200        // Sum should be 15.0
201        assert_eq!(result.data()[0], 15.0);
202    }
203
204    #[test]
205    fn test_mean_cpu_fallback() {
206        let input = Tensor::from_array(array![1.0f32, 2.0, 3.0, 4.0, 5.0].into_dyn());
207        let result = mean_f32_cpu(&input).expect("test: mean_f32_cpu should succeed");
208
209        // Mean should be 3.0
210        assert_eq!(result.data()[0], 3.0);
211    }
212}
tenflowers_core/dispatch_init.rs

tenflowers_core/
dispatch_init.rs
