1use std::collections::HashMap;
6use std::fmt;
7use std::marker::PhantomData;
8use std::sync::{Arc, Mutex};
9
10pub mod async_execution;
11pub mod async_transfer;
12pub mod auto_tuning;
13pub mod backends;
14pub mod benchmarks;
15mod cpu_ops;
16pub mod device_info;
17pub mod heterogeneous;
18pub mod kernels;
19pub mod memory_management;
20pub mod stream_allocator;
21pub mod tensor_cores;
22
23pub use async_transfer::{
24 AsyncTransferError, AsyncTransferPipeline, TransferDirection, TransferHandle,
25};
26pub use device_info::GpuDeviceInfo;
27pub use memory_management::unified_memory::{SyncState, UnifiedAllocator, UnifiedBuffer};
28pub use stream_allocator::{StreamAllocator, StreamId};
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32pub enum GpuBackend {
33 Cuda,
35 Rocm,
37 Wgpu,
39 Metal,
41 OpenCL,
43 Cpu,
45}
46
47impl Default for GpuBackend {
48 fn default() -> Self {
49 Self::preferred()
50 }
51}
52
53impl GpuBackend {
54 pub fn preferred() -> Self {
56 match backends::initialize_optimal_backend() {
59 Ok(backend) => {
60 if backend != GpuBackend::Cpu {
62 #[cfg(not(test))]
65 {
66 return GpuBackend::Cpu;
68 }
69 #[cfg(test)]
70 {
71 return backend;
73 }
74 }
75 backend
76 }
77 Err(_) => {
78 GpuBackend::Cpu
80 }
81 }
82 }
83
84 pub fn is_available(&self) -> bool {
86 match self {
87 GpuBackend::Cuda => false, GpuBackend::Rocm => cfg!(feature = "rocm"), GpuBackend::Wgpu => {
91 #[cfg(feature = "wgpu")]
92 {
93 use crate::gpu::backends::wgpu::WebGPUContext;
94 WebGPUContext::is_available()
95 }
96 #[cfg(not(feature = "wgpu"))]
97 {
98 false
99 }
100 }
101 GpuBackend::Metal => {
102 #[cfg(all(feature = "metal", target_os = "macos"))]
103 {
104 true
106 }
107 #[cfg(not(all(feature = "metal", target_os = "macos")))]
108 {
109 false
110 }
111 }
112 GpuBackend::OpenCL => {
113 #[cfg(feature = "opencl")]
114 {
115 use crate::gpu::backends::opencl::OpenCLContext;
116 OpenCLContext::is_available()
117 }
118 #[cfg(not(feature = "opencl"))]
119 {
120 false
121 }
122 }
123 GpuBackend::Cpu => true,
124 }
125 }
126}
127
128impl fmt::Display for GpuBackend {
129 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130 match self {
131 GpuBackend::Cuda => write!(f, "CUDA"),
132 GpuBackend::Rocm => write!(f, "ROCm"),
133 GpuBackend::Wgpu => write!(f, "WebGPU"),
134 GpuBackend::Metal => write!(f, "Metal"),
135 GpuBackend::OpenCL => write!(f, "OpenCL"),
136 GpuBackend::Cpu => write!(f, "CPU"),
137 }
138 }
139}
140
141use crate::error::{CoreError, ErrorContext, ErrorLocation};
142
143#[derive(Debug, thiserror::Error)]
145pub enum GpuError {
146 #[error("GPU backend {0} is not available")]
148 BackendNotAvailable(String),
149
150 #[error("GPU backend {0} is not supported")]
152 UnsupportedBackend(GpuBackend),
153
154 #[error("GPU backend {0:?} is not supported for this kernel")]
156 BackendNotSupported(GpuBackend),
157
158 #[error("GPU backend {0} is not implemented yet")]
160 BackendNotImplemented(GpuBackend),
161
162 #[error("GPU out of memory: {0}")]
164 OutOfMemory(String),
165
166 #[error("Kernel compilation error: {0}")]
168 KernelCompilationError(String),
169
170 #[error("Kernel execution error: {0}")]
172 KernelExecutionError(String),
173
174 #[error("Invalid parameter: {0}")]
176 InvalidParameter(String),
177
178 #[error("Kernel not found: {0}")]
180 KernelNotFound(String),
181
182 #[error("Kernel specialization not supported")]
184 SpecializationNotSupported,
185
186 #[error("Unsupported data type: {0:?}")]
188 UnsupportedDataType(kernels::DataType),
189
190 #[error("{0}")]
192 Other(String),
193}
194
195#[derive(Debug, Clone, Copy, PartialEq, Eq)]
197pub struct GpuDevice {
198 backend: GpuBackend,
199 device_id: usize,
200}
201
202impl GpuDevice {
203 pub fn new(backend: GpuBackend, device_id: usize) -> Self {
205 Self { backend, device_id }
206 }
207
208 pub fn backend(&self) -> GpuBackend {
210 self.backend
211 }
212
213 pub fn device_id(&self) -> usize {
215 self.device_id
216 }
217
218 pub fn compile_kernel(&self, _source: &str, entrypoint: &str) -> Result<GpuKernel, GpuError> {
220 Ok(GpuKernel {
222 backend: self.backend,
223 entry_point: entrypoint.to_string(),
224 })
225 }
226
227 pub fn get_info(&self) -> Result<GpuDeviceInfo, GpuError> {
244 Ok(GpuDeviceInfo::for_backend(self.backend))
245 }
246}
247
248pub struct GpuKernel {
250 backend: GpuBackend,
251 entry_point: String,
252}
253
254impl GpuKernel {
255 pub fn backend(&self) -> GpuBackend {
257 self.backend
258 }
259
260 pub fn entry_point(&self) -> &str {
262 &self.entry_point
263 }
264}
265
266impl From<GpuError> for CoreError {
268 fn from(err: GpuError) -> Self {
269 match err {
270 GpuError::BackendNotAvailable(backend) => CoreError::ComputationError(
271 ErrorContext::new(format!("GPU backend {backend} is not available"))
272 .with_location(ErrorLocation::new(file!(), line!())),
273 ),
274 GpuError::UnsupportedBackend(backend) => CoreError::NotImplementedError(
275 ErrorContext::new(format!("GPU backend {backend} is not supported"))
276 .with_location(ErrorLocation::new(file!(), line!())),
277 ),
278 GpuError::BackendNotSupported(backend) => CoreError::NotImplementedError(
279 ErrorContext::new(format!(
280 "GPU backend {backend:?} is not supported for this kernel"
281 ))
282 .with_location(ErrorLocation::new(file!(), line!())),
283 ),
284 GpuError::BackendNotImplemented(backend) => CoreError::NotImplementedError(
285 ErrorContext::new(format!("GPU backend {backend} is not implemented yet"))
286 .with_location(ErrorLocation::new(file!(), line!())),
287 ),
288 GpuError::OutOfMemory(details) => CoreError::MemoryError(
289 ErrorContext::new(details.to_string())
290 .with_location(ErrorLocation::new(file!(), line!())),
291 ),
292 GpuError::KernelCompilationError(msg) => CoreError::ComputationError(
293 ErrorContext::new(msg.to_string())
294 .with_location(ErrorLocation::new(file!(), line!())),
295 ),
296 GpuError::KernelExecutionError(msg) => CoreError::ComputationError(
297 ErrorContext::new(msg.to_string())
298 .with_location(ErrorLocation::new(file!(), line!())),
299 ),
300 GpuError::InvalidParameter(msg) => CoreError::InvalidArgument(
301 ErrorContext::new(msg.to_string())
302 .with_location(ErrorLocation::new(file!(), line!())),
303 ),
304 GpuError::KernelNotFound(name) => CoreError::ComputationError(
305 ErrorContext::new(name.to_string())
306 .with_location(ErrorLocation::new(file!(), line!())),
307 ),
308 GpuError::SpecializationNotSupported => CoreError::NotImplementedError(
309 ErrorContext::new("Kernel specialization not supported".to_string())
310 .with_location(ErrorLocation::new(file!(), line!())),
311 ),
312 GpuError::UnsupportedDataType(dtype) => CoreError::TypeError(
313 ErrorContext::new(format!("{dtype:?}"))
314 .with_location(ErrorLocation::new(file!(), line!())),
315 ),
316 GpuError::Other(msg) => CoreError::ComputationError(
317 ErrorContext::new(msg).with_location(ErrorLocation::new(file!(), line!())),
318 ),
319 }
320 }
321}
322
323pub trait GpuDataType: Copy + Send + Sync + 'static {}
325
326#[derive(Debug)]
328pub struct GpuPtr<T: GpuDataType> {
329 ptr: u64,
330 size: usize,
331 phantom: PhantomData<T>,
332}
333
334impl<T: GpuDataType> GpuPtr<T> {
335 pub fn allocate(size: usize) -> Result<Self, GpuError> {
337 Ok(GpuPtr {
338 ptr: 0x1000_0000, size,
340 phantom: PhantomData,
341 })
342 }
343
344 pub fn as_ptr(&self) -> u64 {
346 self.ptr
347 }
348
349 pub fn len(&self) -> usize {
351 self.size
352 }
353
354 pub fn is_empty(&self) -> bool {
356 self.size == 0
357 }
358}
359
360#[derive(Debug, Clone)]
362pub enum KernelArg<'a, T: GpuDataType> {
363 Buffer(&'a GpuPtr<T>),
365 Scalar(T),
367}
368
369#[derive(Debug, Clone)]
371pub enum DynamicKernelArg {
372 Buffer(u64), F32(f32),
376 F64(f64),
378 I32(i32),
380 U32(u32),
382 Usize(usize),
384}
385
386pub struct GpuChannel {
388 #[allow(dead_code)]
389 source_device: usize,
390 #[allow(dead_code)]
391 target_device: usize,
392 #[allow(dead_code)]
393 bandwidth: f64, }
395
396impl GpuDataType for f32 {}
398impl GpuDataType for f64 {}
399impl GpuDataType for i32 {}
400impl GpuDataType for u32 {}
401impl GpuDataType for u8 {}
402impl GpuDataType for i8 {}
403impl GpuDataType for u16 {}
404impl GpuDataType for i16 {}
405impl GpuDataType for u64 {}
406impl GpuDataType for i64 {}
407impl GpuDataType for usize {}
408impl GpuDataType for isize {}
409
410pub struct GpuBuffer<T: GpuDataType> {
412 inner: Arc<dyn GpuBufferImpl>,
413 size: usize,
414 phantom: PhantomData<T>,
415}
416
417impl<T: GpuDataType> GpuBuffer<T> {
418 pub(crate) fn new(inner: Arc<dyn GpuBufferImpl>, size: usize) -> Self {
420 Self {
421 inner,
422 size,
423 phantom: PhantomData,
424 }
425 }
426
427 pub fn len(&self) -> usize {
429 self.size
430 }
431
432 pub fn is_empty(&self) -> bool {
434 self.size == 0
435 }
436
437 pub fn copy_from_host(&self, data: &[T]) -> Result<(), GpuError> {
439 if data.len() > self.size {
440 return Err(GpuError::InvalidParameter(
441 "Data size exceeds buffer size".to_string(),
442 ));
443 }
444 unsafe {
445 self.inner
446 .copy_from_host(data.as_ptr() as *const u8, std::mem::size_of_val(data));
447 }
448 Ok(())
449 }
450
451 pub fn copy_to_host(&self, data: &mut [T]) -> Result<(), GpuError> {
453 if data.len() > self.size {
454 return Err(GpuError::InvalidParameter(
455 "Data size exceeds buffer size".to_string(),
456 ));
457 }
458 unsafe {
459 self.inner
460 .copy_to_host(data.as_mut_ptr() as *mut u8, std::mem::size_of_val(data));
461 }
462 Ok(())
463 }
464
465 pub fn to_vec(&self) -> Vec<T> {
467 let mut result = vec![unsafe { std::mem::zeroed() }; self.size];
468 let _ = self.copy_to_host(&mut result);
469 result
470 }
471}
472
473impl<T: GpuDataType> fmt::Debug for GpuBuffer<T> {
474 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
475 f.debug_struct("GpuBuffer")
476 .field("size", &self.size)
477 .finish()
478 }
479}
480
481impl<T: GpuDataType> Clone for GpuBuffer<T> {
482 fn clone(&self) -> Self {
483 Self {
484 inner: Arc::clone(&self.inner),
485 size: self.size,
486 phantom: PhantomData,
487 }
488 }
489}
490
491#[derive(Clone)]
493pub struct GpuKernelHandle {
494 inner: Arc<dyn GpuKernelImpl>,
495}
496
497impl GpuKernelHandle {
498 pub(crate) fn new(inner: Arc<dyn GpuKernelImpl>) -> Self {
500 Self { inner }
501 }
502
503 pub fn set_buffer<T: GpuDataType>(&self, name: &str, buffer: &GpuBuffer<T>) {
505 self.inner.set_buffer(name, &buffer.inner);
506 }
507
508 pub fn set_u32(&self, name: &str, value: u32) {
510 self.inner.set_u32(name, value);
511 }
512
513 pub fn set_i32(&self, name: &str, value: i32) {
515 self.inner.set_i32(name, value);
516 }
517
518 pub fn set_f32(&self, name: &str, value: f32) {
520 self.inner.set_f32(name, value);
521 }
522
523 pub fn set_f64(&self, name: &str, value: f64) {
525 self.inner.set_f64(name, value);
526 }
527
528 pub fn dispatch(&self, workgroups: [u32; 3]) {
534 if !self.inner.try_batch_dispatch(workgroups) {
535 self.inner.dispatch(workgroups);
536 }
537 }
538
539 pub fn dispatch_no_wait(&self, workgroups: [u32; 3]) {
544 self.inner.dispatch_no_wait(workgroups);
545 }
546}
547
548pub struct GpuCompiler {
550 inner: Arc<dyn GpuCompilerImpl>,
551}
552
553impl GpuCompiler {
554 pub(crate) fn new(inner: Arc<dyn GpuCompilerImpl>) -> Self {
556 Self { inner }
557 }
558
559 pub fn compile(&self, source: &str) -> Result<GpuKernelHandle, GpuError> {
561 let kernel = self.inner.compile(source)?;
562 Ok(GpuKernelHandle::new(kernel))
563 }
564
565 pub fn compile_kernel<I: GpuDataType, O: GpuDataType>(
575 &self,
576 name: &str,
577 ) -> Result<GpuKernelHandle, GpuError> {
578 let kernel = self.inner.compile_typed(
579 name,
580 std::any::TypeId::of::<I>(),
581 std::any::TypeId::of::<O>(),
582 )?;
583 Ok(GpuKernelHandle::new(kernel))
584 }
585}
586
587pub struct GpuContext {
589 inner: Arc<dyn GpuContextImpl>,
590 backend: GpuBackend,
591 kernel_registry: kernels::KernelRegistry,
592}
593
594impl GpuContext {
595 pub fn new(backend: GpuBackend) -> Result<Self, GpuError> {
597 if backend == GpuBackend::Cuda {
601 return Err(GpuError::BackendNotAvailable(
602 "CUDA backend was removed from scirs2-core in 0.6.x; use the oxicuda-* crates directly (per-crate `cuda` feature). See SCIRS2_POLICY.md".to_string(),
603 ));
604 }
605
606 if !backend.is_available() {
608 return Err(GpuError::BackendNotAvailable(backend.to_string()));
609 }
610
611 if backend != GpuBackend::Cpu {
613 let detection_result = backends::detect_gpu_backends();
614 let backend_available = detection_result
615 .devices
616 .iter()
617 .any(|d| d.backend == backend && d.backend != GpuBackend::Cpu);
618
619 if !backend_available {
620 return Err(GpuError::BackendNotAvailable(format!(
621 "{backend} (no devices detected at runtime)"
622 )));
623 }
624 }
625
626 let inner = match backend {
627 GpuBackend::Cuda => {
628 return Err(GpuError::BackendNotAvailable(
629 "CUDA backend was removed from scirs2-core in 0.6.x; use the oxicuda-* crates directly (per-crate `cuda` feature). See SCIRS2_POLICY.md".to_string(),
630 ));
631 }
632 GpuBackend::Rocm => {
633 #[cfg(feature = "rocm")]
634 {
635 #[cfg(test)]
638 {
639 Arc::new(CpuContext::new()) as Arc<dyn GpuContextImpl>
641 }
642 #[cfg(not(test))]
643 {
644 return Err(GpuError::BackendNotImplemented(backend));
645 }
646 }
647 #[cfg(not(feature = "rocm"))]
648 {
649 return Err(GpuError::UnsupportedBackend(backend));
650 }
651 }
652 GpuBackend::Wgpu => {
653 #[cfg(feature = "wgpu")]
654 {
655 use crate::gpu::backends::wgpu::WebGPUContext;
656 match WebGPUContext::new() {
657 Ok(ctx) => Arc::new(ctx) as Arc<dyn GpuContextImpl>,
658 Err(e) => return Err(e),
659 }
660 }
661 #[cfg(not(feature = "wgpu"))]
662 {
663 return Err(GpuError::UnsupportedBackend(backend));
664 }
665 }
666 GpuBackend::Metal => {
667 #[cfg(all(feature = "metal", target_os = "macos"))]
668 {
669 use crate::gpu::backends::metal::MetalContext;
670 match MetalContext::new() {
671 Ok(ctx) => Arc::new(ctx) as Arc<dyn GpuContextImpl>,
672 Err(e) => return Err(e),
673 }
674 }
675 #[cfg(not(all(feature = "metal", target_os = "macos")))]
676 {
677 return Err(GpuError::UnsupportedBackend(backend));
678 }
679 }
680 GpuBackend::OpenCL => {
681 #[cfg(feature = "opencl")]
682 {
683 use crate::gpu::backends::opencl::OpenCLContext;
684 match OpenCLContext::new() {
685 Ok(ctx) => Arc::new(ctx) as Arc<dyn GpuContextImpl>,
686 Err(e) => return Err(e),
687 }
688 }
689 #[cfg(not(feature = "opencl"))]
690 {
691 return Err(GpuError::UnsupportedBackend(backend));
692 }
693 }
694 GpuBackend::Cpu => Arc::new(CpuContext::new()) as Arc<dyn GpuContextImpl>,
695 };
696
697 Ok(Self {
698 inner,
699 backend,
700 kernel_registry: kernels::KernelRegistry::with_default_kernels(),
701 })
702 }
703
704 pub fn backend(&self) -> GpuBackend {
706 self.backend
707 }
708
709 pub fn backend_name(&self) -> &str {
711 match self.backend {
712 GpuBackend::Cuda => "CUDA",
713 GpuBackend::Rocm => "ROCm",
714 GpuBackend::Wgpu => "WebGPU",
715 GpuBackend::Metal => "Metal",
716 GpuBackend::OpenCL => "OpenCL",
717 GpuBackend::Cpu => "CPU",
718 }
719 }
720
721 pub fn gpu_sync(&self) -> Result<(), GpuError> {
726 self.inner.gpu_sync()
727 }
728
729 pub fn begin_batch(&self) -> Result<(), GpuError> {
735 self.inner.begin_batch()
736 }
737
738 pub fn end_batch(&self) -> Result<(), GpuError> {
744 self.inner.end_batch()
745 }
746
747 pub fn create_buffer<T: GpuDataType>(&self, size: usize) -> GpuBuffer<T> {
748 let byte_size = size.saturating_mul(std::mem::size_of::<T>());
749 let inner = self.inner.create_buffer(byte_size);
750 GpuBuffer::new(inner, size)
751 }
752
753 pub fn create_buffer_from_slice<T: GpuDataType>(&self, data: &[T]) -> GpuBuffer<T> {
755 let buffer = self.create_buffer::<T>(data.len());
756 let _ = buffer.copy_from_host(data);
757 buffer
758 }
759
760 pub fn execute<F, R>(&self, f: F) -> R
762 where
763 F: FnOnce(&GpuCompiler) -> R,
764 {
765 let compiler = GpuCompiler::new(self.inner.create_compiler());
766 f(&compiler)
767 }
768
769 pub fn get_kernel(&self, name: &str) -> Result<GpuKernelHandle, GpuError> {
771 let kernel = self
772 .kernel_registry
773 .get(name)
774 .ok_or_else(|| GpuError::KernelNotFound(name.to_string()))?;
775
776 let kernel_source = kernel.source_for_backend(self.backend)?;
777 let metadata = kernel.metadata();
778
779 let handle = self.compile_kernel_with_metadata(&kernel_source, &metadata)?;
780 Ok(handle)
781 }
782
783 pub fn get_specialized_kernel(
785 &self,
786 name: &str,
787 params: &kernels::KernelParams,
788 ) -> Result<GpuKernelHandle, GpuError> {
789 let specialized = self.kernel_registry.get_specialized(name, params)?;
790 let kernel_source = specialized.source_for_backend(self.backend)?;
791 let metadata = specialized.metadata();
792
793 let handle = self.compile_kernel_with_metadata(&kernel_source, &metadata)?;
794 Ok(handle)
795 }
796
797 fn compile_kernel_with_metadata(
799 &self,
800 source: &str,
801 _metadata: &kernels::KernelMetadata,
802 ) -> Result<GpuKernelHandle, GpuError> {
803 self.execute(|compiler| compiler.compile(source))
804 }
805
806 pub fn get_available_memory(&self) -> Option<usize> {
808 Some(1024 * 1024 * 1024) }
812
813 pub fn get_total_memory(&self) -> Option<usize> {
815 #[cfg(target_arch = "wasm32")]
818 return Some(512 * 1024 * 1024); #[cfg(not(target_arch = "wasm32"))]
821 Some((4u64 * 1024 * 1024 * 1024) as usize) }
823
824 pub fn launch_kernel(
826 &self,
827 kernel_name: &str,
828 grid_size: (usize, usize, usize),
829 block_size: (usize, usize, usize),
830 args: &[DynamicKernelArg],
831 ) -> Result<(), GpuError> {
832 let _ = (kernel_name, grid_size, block_size, args);
834 Ok(())
835 }
836
837 pub fn transfer_async_host_to_device<T: GpuDataType>(
839 &self,
840 ptr: &GpuPtr<T>,
841 data: &[T],
842 ) -> Result<(), GpuError> {
843 let _ = (ptr, data);
845 Ok(())
846 }
847
848 pub fn transfer_host_to_device<T: GpuDataType>(
850 &self,
851 ptr: &GpuPtr<T>,
852 data: &[T],
853 ) -> Result<(), GpuError> {
854 let _ = (ptr, data);
856 Ok(())
857 }
858
859 pub fn transfer_async_device_to_host<T: GpuDataType>(
861 &self,
862 ptr: &GpuPtr<T>,
863 data: &mut [T],
864 ) -> Result<(), GpuError> {
865 let _ = (ptr, data);
867 Ok(())
868 }
869
870 pub fn transfer_device_to_host<T: GpuDataType>(
872 &self,
873 ptr: &GpuPtr<T>,
874 data: &mut [T],
875 ) -> Result<(), GpuError> {
876 let _ = (ptr, data);
878 Ok(())
879 }
880
881 pub fn execute_kernel(
884 &self,
885 source: &str,
886 buffers: &[GpuBuffer<f32>],
887 work_groups: (u32, u32, u32),
888 int_params: &[u32],
889 float_params: &[f32],
890 ) -> Result<(), GpuError> {
891 eprintln!(
894 "GPU kernel execution (source length: {}, buffers: {}, workgroups: {:?})",
895 source.len(),
896 buffers.len(),
897 work_groups
898 );
899 eprintln!("Int params: {int_params:?}");
900 eprintln!("Float params: {float_params:?}");
901 Ok(())
902 }
903
904 pub fn read_buffer<T: GpuDataType>(&self, buffer: &GpuBuffer<T>) -> Result<Vec<T>, GpuError> {
907 Ok(buffer.to_vec())
908 }
909
910 pub fn sum_all<T: GpuDataType>(&self, buffer: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
912 self.sum_all_cpu_fallback(buffer)
913 }
914
915 pub fn mean_all<T: GpuDataType>(
917 &self,
918 buffer: &GpuBuffer<T>,
919 ) -> Result<GpuBuffer<T>, GpuError> {
920 self.mean_all_cpu_fallback(buffer)
921 }
922
923 pub fn max_all<T: GpuDataType>(&self, buffer: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
925 self.max_all_cpu_fallback(buffer)
926 }
927
928 pub fn min_all<T: GpuDataType>(&self, buffer: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
930 self.min_all_cpu_fallback(buffer)
931 }
932
933 pub fn sum_axis<T: GpuDataType>(
935 &self,
936 buffer: &GpuBuffer<T>,
937 shape: &[usize],
938 axis: usize,
939 ) -> Result<GpuBuffer<T>, GpuError> {
940 self.sum_axis_cpu_fallback(buffer, shape, axis)
941 }
942
943 pub fn mean_axis<T: GpuDataType>(
945 &self,
946 buffer: &GpuBuffer<T>,
947 shape: &[usize],
948 axis: usize,
949 ) -> Result<GpuBuffer<T>, GpuError> {
950 self.mean_axis_cpu_fallback(buffer, shape, axis)
951 }
952
953 pub fn max_axis<T: GpuDataType>(
955 &self,
956 buffer: &GpuBuffer<T>,
957 shape: &[usize],
958 axis: usize,
959 ) -> Result<GpuBuffer<T>, GpuError> {
960 self.max_axis_cpu_fallback(buffer, shape, axis)
961 }
962
963 pub fn min_axis<T: GpuDataType>(
965 &self,
966 buffer: &GpuBuffer<T>,
967 shape: &[usize],
968 axis: usize,
969 ) -> Result<GpuBuffer<T>, GpuError> {
970 self.min_axis_cpu_fallback(buffer, shape, axis)
971 }
972
973 pub fn broadcast<T: GpuDataType>(
975 &self,
976 buffer: &GpuBuffer<T>,
977 from_shape: &[usize],
978 to_shape: &[usize],
979 ) -> Result<GpuBuffer<T>, GpuError> {
980 self.broadcast_cpu_fallback(buffer, from_shape, to_shape)
981 }
982
983 pub fn scale<T: GpuDataType>(
985 &self,
986 buffer: &GpuBuffer<T>,
987 scalar: T,
988 ) -> Result<GpuBuffer<T>, GpuError> {
989 self.scale_cpu_fallback(buffer, scalar)
990 }
991
992 pub fn gemm<T: GpuDataType>(
994 &self,
995 a: &GpuBuffer<T>,
996 b: &GpuBuffer<T>,
997 m: usize,
998 k: usize,
999 n: usize,
1000 ) -> Result<GpuBuffer<T>, GpuError> {
1001 self.gemm_cpu_fallback(a, b, m, k, n)
1002 }
1003
1004 pub fn gemm_transpose_b<T: GpuDataType>(
1006 &self,
1007 a: &GpuBuffer<T>,
1008 b: &GpuBuffer<T>,
1009 m: usize,
1010 k: usize,
1011 n: usize,
1012 ) -> Result<GpuBuffer<T>, GpuError> {
1013 self.gemm_transpose_b_cpu_fallback(a, b, m, k, n)
1014 }
1015
1016 pub fn gemm_transpose_a<T: GpuDataType>(
1018 &self,
1019 a: &GpuBuffer<T>,
1020 b: &GpuBuffer<T>,
1021 m: usize,
1022 k: usize,
1023 n: usize,
1024 ) -> Result<GpuBuffer<T>, GpuError> {
1025 self.gemm_transpose_a_cpu_fallback(a, b, m, k, n)
1026 }
1027
1028 pub fn relu<T: GpuDataType>(&self, input: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
1030 self.relu_cpu_fallback(input)
1031 }
1032
1033 pub fn relu_backward<T: GpuDataType>(
1035 &self,
1036 grad_output: &GpuBuffer<T>,
1037 input: &GpuBuffer<T>,
1038 ) -> Result<GpuBuffer<T>, GpuError> {
1039 self.relu_backward_cpu_fallback(grad_output, input)
1040 }
1041
1042 pub fn sigmoid<T: GpuDataType>(&self, input: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
1044 self.sigmoid_cpu_fallback(input)
1045 }
1046
1047 pub fn sigmoid_backward<T: GpuDataType>(
1049 &self,
1050 grad_output: &GpuBuffer<T>,
1051 input: &GpuBuffer<T>,
1052 ) -> Result<GpuBuffer<T>, GpuError> {
1053 self.sigmoid_backward_cpu_fallback(grad_output, input)
1054 }
1055
1056 pub fn tanh<T: GpuDataType>(&self, input: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
1058 self.tanh_cpu_fallback(input)
1059 }
1060
1061 pub fn tanh_backward<T: GpuDataType>(
1063 &self,
1064 grad_output: &GpuBuffer<T>,
1065 input: &GpuBuffer<T>,
1066 ) -> Result<GpuBuffer<T>, GpuError> {
1067 self.tanh_backward_cpu_fallback(grad_output, input)
1068 }
1069
1070 pub fn gelu<T: GpuDataType>(&self, input: &GpuBuffer<T>) -> Result<GpuBuffer<T>, GpuError> {
1072 self.gelu_cpu_fallback(input)
1073 }
1074
1075 pub fn gelu_backward<T: GpuDataType>(
1077 &self,
1078 grad_output: &GpuBuffer<T>,
1079 input: &GpuBuffer<T>,
1080 ) -> Result<GpuBuffer<T>, GpuError> {
1081 self.gelu_backward_cpu_fallback(grad_output, input)
1082 }
1083}
1084
1085impl fmt::Debug for GpuContext {
1086 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1087 f.debug_struct("GpuContext")
1088 .field("backend", &self.backend)
1089 .finish()
1090 }
1091}
1092
1093pub(crate) trait GpuBufferImpl: Send + Sync {
1098 unsafe fn copy_from_host(&self, data: *const u8, size: usize);
1100
1101 unsafe fn copy_to_host(&self, data: *mut u8, size: usize);
1103
1104 #[allow(dead_code)]
1106 fn as_any(&self) -> &dyn std::any::Any;
1107
1108 #[allow(dead_code)]
1110 fn size(&self) -> usize {
1111 0 }
1113
1114 #[allow(dead_code)]
1116 fn device_ptr(&self) -> u64 {
1117 0 }
1119}
1120
1121pub(crate) trait GpuKernelImpl: Send + Sync {
1123 fn set_buffer(&self, name: &str, buffer: &Arc<dyn GpuBufferImpl>);
1125
1126 fn set_u32(&self, name: &str, value: u32);
1128
1129 fn set_i32(&self, name: &str, value: i32);
1131
1132 fn set_f32(&self, name: &str, value: f32);
1134
1135 fn set_f64(&self, name: &str, value: f64);
1137
1138 fn dispatch(&self, workgroups: [u32; 3]);
1140
1141 fn dispatch_no_wait(&self, workgroups: [u32; 3]) {
1145 self.dispatch(workgroups);
1147 }
1148
1149 fn try_batch_dispatch(&self, _workgroups: [u32; 3]) -> bool {
1154 false
1155 }
1156}
1157
1158pub(crate) trait GpuCompilerImpl: Send + Sync {
1160 fn compile(&self, source: &str) -> Result<Arc<dyn GpuKernelImpl>, GpuError>;
1162
1163 fn compile_typed(
1169 &self,
1170 name: &str,
1171 input_type: std::any::TypeId,
1172 output_type: std::any::TypeId,
1173 ) -> Result<Arc<dyn GpuKernelImpl>, GpuError>;
1174}
1175
1176pub(crate) trait GpuContextImpl: Send + Sync {
1178 fn create_buffer(&self, size: usize) -> Arc<dyn GpuBufferImpl>;
1180
1181 fn create_compiler(&self) -> Arc<dyn GpuCompilerImpl>;
1183
1184 fn gpu_sync(&self) -> Result<(), GpuError> {
1186 Ok(()) }
1188
1189 fn begin_batch(&self) -> Result<(), GpuError> {
1194 Ok(()) }
1196
1197 fn end_batch(&self) -> Result<(), GpuError> {
1202 Ok(()) }
1204
1205 fn as_any(&self) -> &dyn std::any::Any
1207 where
1208 Self: 'static + Sized,
1209 {
1210 self
1211 }
1212}
1213
1214struct CpuContext;
1218
1219impl CpuContext {
1220 fn new() -> Self {
1222 Self
1223 }
1224}
1225
1226impl GpuContextImpl for CpuContext {
1227 fn create_buffer(&self, size: usize) -> Arc<dyn GpuBufferImpl> {
1228 Arc::new(CpuBuffer::new(size))
1229 }
1230
1231 fn create_compiler(&self) -> Arc<dyn GpuCompilerImpl> {
1232 Arc::new(CpuCompiler)
1233 }
1234}
1235
1236struct CpuBuffer {
1238 data: Vec<u8>,
1239}
1240
1241impl CpuBuffer {
1242 fn new(size: usize) -> Self {
1244 Self {
1245 data: vec![0; size],
1246 }
1247 }
1248}
1249
1250impl GpuBufferImpl for CpuBuffer {
1251 unsafe fn copy_from_host(&self, data: *const u8, size: usize) {
1252 let mut_self = self as *const Self as *mut Self;
1253 let data_ptr = (*mut_self).data.as_mut_ptr();
1254 std::ptr::copy_nonoverlapping(data, data_ptr, size);
1255 }
1256
1257 unsafe fn copy_to_host(&self, data: *mut u8, size: usize) {
1258 let data_ptr = self.data.as_ptr();
1259 std::ptr::copy_nonoverlapping(data_ptr, data, size);
1260 }
1261
1262 fn as_any(&self) -> &dyn std::any::Any {
1263 self
1264 }
1265
1266 fn size(&self) -> usize {
1267 self.data.len()
1268 }
1269
1270 fn device_ptr(&self) -> u64 {
1271 self.data.as_ptr() as u64
1272 }
1273}
1274
1275struct CpuCompiler;
1277
1278impl GpuCompilerImpl for CpuCompiler {
1279 fn compile(&self, source: &str) -> Result<Arc<dyn GpuKernelImpl>, GpuError> {
1280 Ok(Arc::new(CpuKernel))
1283 }
1284
1285 fn compile_typed(
1286 &self,
1287 name: &str,
1288 input_type: std::any::TypeId,
1289 output_type: std::any::TypeId,
1290 ) -> Result<Arc<dyn GpuKernelImpl>, GpuError> {
1291 let is_f32 = input_type == std::any::TypeId::of::<f32>()
1298 && output_type == std::any::TypeId::of::<f32>();
1299 if name == "vector_add" && is_f32 {
1300 return Ok(Arc::new(CpuVectorAddKernel::new()));
1301 }
1302
1303 Err(GpuError::KernelCompilationError(format!(
1304 "no CPU fallback implementation is registered for kernel '{name}' with the \
1305 requested types; only a real f32 'vector_add' kernel is available on the \
1306 pure-Rust CPU backend without a compiled GPU shader. Use GpuCompiler::compile \
1307 with real kernel source for other kernels."
1308 )))
1309 }
1310}
1311
1312struct CpuKernel;
1314
1315impl GpuKernelImpl for CpuKernel {
1316 fn set_buffer(&self, _name: &str, buffer: &Arc<dyn GpuBufferImpl>) {
1317 }
1319
1320 fn set_u32(&self, _name: &str, value: u32) {
1321 }
1323
1324 fn set_i32(&self, _name: &str, value: i32) {
1325 }
1327
1328 fn set_f32(&self, _name: &str, value: f32) {
1329 }
1331
1332 fn set_f64(&self, _name: &str, value: f64) {
1333 }
1335
1336 fn dispatch(&self, workgroups: [u32; 3]) {
1337 }
1339}
1340
1341struct CpuVectorAddKernel {
1348 buffers: Mutex<HashMap<String, Arc<dyn GpuBufferImpl>>>,
1349}
1350
1351impl CpuVectorAddKernel {
1352 fn new() -> Self {
1353 Self {
1354 buffers: Mutex::new(HashMap::new()),
1355 }
1356 }
1357}
1358
1359impl GpuKernelImpl for CpuVectorAddKernel {
1360 fn set_buffer(&self, name: &str, buffer: &Arc<dyn GpuBufferImpl>) {
1361 if let Ok(mut buffers) = self.buffers.lock() {
1362 buffers.insert(name.to_string(), Arc::clone(buffer));
1363 }
1364 }
1365
1366 fn set_u32(&self, _name: &str, _value: u32) {}
1367 fn set_i32(&self, _name: &str, _value: i32) {}
1368 fn set_f32(&self, _name: &str, _value: f32) {}
1369 fn set_f64(&self, _name: &str, _value: f64) {}
1370
1371 fn dispatch(&self, _workgroups: [u32; 3]) {
1372 let Ok(buffers) = self.buffers.lock() else {
1373 eprintln!("Warning: CPU vector_add dispatch: buffer registry lock poisoned");
1374 return;
1375 };
1376 let (Some(a), Some(b), Some(result)) =
1377 (buffers.get("a"), buffers.get("b"), buffers.get("result"))
1378 else {
1379 eprintln!(
1380 "Warning: CPU vector_add dispatch requires buffers named 'a', 'b', and \
1381 'result' (set via set_buffer); dispatch skipped"
1382 );
1383 return;
1384 };
1385
1386 let byte_len = result.size();
1387 if a.size() < byte_len || b.size() < byte_len {
1388 eprintln!(
1389 "Warning: CPU vector_add dispatch: input buffers ({}, {} bytes) are smaller \
1390 than the output buffer ({byte_len} bytes); dispatch skipped",
1391 a.size(),
1392 b.size()
1393 );
1394 return;
1395 }
1396
1397 let mut a_bytes = vec![0u8; byte_len];
1398 let mut b_bytes = vec![0u8; byte_len];
1399 unsafe {
1403 a.copy_to_host(a_bytes.as_mut_ptr(), byte_len);
1404 b.copy_to_host(b_bytes.as_mut_ptr(), byte_len);
1405 }
1406
1407 let read_f32 = |bytes: &[u8]| -> Vec<f32> {
1408 bytes
1409 .chunks_exact(4)
1410 .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
1411 .collect()
1412 };
1413 let a_vals = read_f32(&a_bytes);
1414 let b_vals = read_f32(&b_bytes);
1415
1416 let result_bytes: Vec<u8> = a_vals
1417 .iter()
1418 .zip(b_vals.iter())
1419 .flat_map(|(x, y)| (x + y).to_le_bytes())
1420 .collect();
1421
1422 unsafe {
1424 result.copy_from_host(result_bytes.as_ptr(), result_bytes.len());
1425 }
1426 }
1427}
1428
1429#[cfg(test)]
1433mod tests {
1434 use super::*;
1435
1436 #[test]
1437 fn test_gpu_backend_preferred() {
1438 let backend = GpuBackend::preferred();
1439 match backend {
1441 GpuBackend::Cuda
1442 | GpuBackend::Rocm
1443 | GpuBackend::Wgpu
1444 | GpuBackend::Metal
1445 | GpuBackend::OpenCL
1446 | GpuBackend::Cpu => {}
1447 }
1448 }
1449
1450 #[test]
1451 fn test_gpu_backend_default() {
1452 let backend = GpuBackend::default();
1453 assert_eq!(backend, GpuBackend::preferred());
1454 }
1455
1456 #[test]
1457 fn test_gpu_backend_is_available() {
1458 let backend = GpuBackend::Cpu;
1459 assert!(backend.is_available());
1460
1461 assert!(!GpuBackend::Cuda.is_available());
1464
1465 #[cfg(feature = "rocm")]
1466 {
1467 let _ = GpuBackend::Rocm.is_available(); }
1470 #[cfg(not(feature = "rocm"))]
1471 assert!(!GpuBackend::Rocm.is_available());
1472
1473 #[cfg(all(feature = "metal", target_os = "macos"))]
1474 assert!(GpuBackend::Metal.is_available());
1475 #[cfg(not(all(feature = "metal", target_os = "macos")))]
1476 assert!(!GpuBackend::Metal.is_available());
1477 }
1478
1479 #[test]
1480 fn test_gpu_backend_display() {
1481 assert_eq!(GpuBackend::Cuda.to_string(), "CUDA");
1482 assert_eq!(GpuBackend::Rocm.to_string(), "ROCm");
1483 assert_eq!(GpuBackend::Wgpu.to_string(), "WebGPU");
1484 assert_eq!(GpuBackend::Metal.to_string(), "Metal");
1485 assert_eq!(GpuBackend::OpenCL.to_string(), "OpenCL");
1486 assert_eq!(GpuBackend::Cpu.to_string(), "CPU");
1487 }
1488
1489 #[test]
1490 fn test_gpuerror_from_conversion() {
1491 let gpuerror = GpuError::BackendNotAvailable("CUDA".to_string());
1492 let coreerror: CoreError = gpuerror.into();
1493 match coreerror {
1494 CoreError::ComputationError(_) => {}
1495 _ => panic!("Expected ComputationError"),
1496 }
1497
1498 let gpuerror = GpuError::OutOfMemory("8GB required".to_string());
1499 let coreerror: CoreError = gpuerror.into();
1500 match coreerror {
1501 CoreError::MemoryError(_) => {}
1502 _ => panic!("Expected MemoryError"),
1503 }
1504
1505 let gpuerror = GpuError::InvalidParameter("batch_size must be > 0".to_string());
1506 let coreerror: CoreError = gpuerror.into();
1507 match coreerror {
1508 CoreError::InvalidArgument(_) => {}
1509 _ => panic!("Expected InvalidArgument"),
1510 }
1511
1512 let gpuerror = GpuError::UnsupportedDataType(kernels::DataType::Float16);
1513 let coreerror: CoreError = gpuerror.into();
1514 match coreerror {
1515 CoreError::TypeError(_) => {}
1516 _ => panic!("Expected TypeError"),
1517 }
1518 }
1519
1520 #[test]
1521 fn test_gpu_datatype_trait() {
1522 fn assert_gpu_datatype<T: GpuDataType>() {}
1524
1525 assert_gpu_datatype::<f32>();
1526 assert_gpu_datatype::<f64>();
1527 assert_gpu_datatype::<i32>();
1528 assert_gpu_datatype::<u32>();
1529 assert_gpu_datatype::<u8>();
1530 assert_gpu_datatype::<i8>();
1531 assert_gpu_datatype::<u16>();
1532 assert_gpu_datatype::<i16>();
1533 assert_gpu_datatype::<u64>();
1534 assert_gpu_datatype::<i64>();
1535 }
1536
1537 #[test]
1538 fn test_gpu_buffer_creation() {
1539 let inner = Arc::new(CpuBuffer::new(100));
1540 let buffer = GpuBuffer::<f32>::new(inner, 25);
1541
1542 assert_eq!(buffer.len(), 25);
1543 assert!(!buffer.is_empty());
1544 }
1545
1546 #[test]
1547 fn test_gpu_buffer_empty() {
1548 let inner = Arc::new(CpuBuffer::new(0));
1549 let buffer = GpuBuffer::<f32>::new(inner, 0);
1550
1551 assert_eq!(buffer.len(), 0);
1552 assert!(buffer.is_empty());
1553 }
1554
1555 #[test]
1556 fn test_gpu_buffer_copy_operations() {
1557 let inner = Arc::new(CpuBuffer::new(16));
1558 let buffer = GpuBuffer::<f32>::new(inner, 4);
1559
1560 let data = vec![1.0f32, 2.0, 3.0, 4.0];
1561 let _ = buffer.copy_from_host(&data);
1562
1563 let mut result = vec![0.0f32; 4];
1564 let _ = buffer.copy_to_host(&mut result);
1565
1566 assert_eq!(result, data);
1567 }
1568
1569 #[test]
1570 fn test_gpu_buffer_to_vec() {
1571 let inner = Arc::new(CpuBuffer::new(12));
1572 let buffer = GpuBuffer::<f32>::new(inner, 3);
1573
1574 let data = vec![5.0f32, 6.0, 7.0];
1575 let _ = buffer.copy_from_host(&data);
1576
1577 let result = buffer.to_vec();
1578 assert_eq!(result, data);
1579 }
1580
1581 #[test]
1582 #[should_panic(expected = "Data size exceeds buffer size")]
1583 fn test_gpu_buffer_copy_from_host_overflow() {
1584 let inner = Arc::new(CpuBuffer::new(8));
1585 let buffer = GpuBuffer::<f32>::new(inner, 2);
1586
1587 let data = vec![1.0f32, 2.0, 3.0]; buffer.copy_from_host(&data).expect("Operation failed");
1589 }
1590
1591 #[test]
1592 #[should_panic(expected = "Data size exceeds buffer size")]
1593 fn test_gpu_buffer_copy_to_host_overflow() {
1594 let inner = Arc::new(CpuBuffer::new(8));
1595 let buffer = GpuBuffer::<f32>::new(inner, 2);
1596
1597 let mut data = vec![0.0f32; 3]; buffer.copy_to_host(&mut data).expect("Operation failed");
1599 }
1600
1601 #[test]
1602 fn test_gpu_kernel_handle() {
1603 let kernel = Arc::new(CpuKernel);
1604 let handle = GpuKernelHandle::new(kernel);
1605
1606 let buffer = GpuBuffer::<f32>::new(Arc::new(CpuBuffer::new(16)), 4);
1608 handle.set_buffer("input", &buffer);
1609 handle.set_u32("size", 100);
1610 handle.set_i32("offset", -5);
1611 handle.set_f32("scale", 2.5);
1612 handle.set_f64("precision", 0.0001);
1613
1614 handle.dispatch([16, 8, 1]);
1616 }
1617
1618 #[test]
1619 fn test_gpu_context_cpu_backend() {
1620 let context = GpuContext::new(GpuBackend::Cpu).expect("Operation failed");
1621 assert_eq!(context.backend(), GpuBackend::Cpu);
1622 assert_eq!(context.backend_name(), "CPU");
1623
1624 assert_eq!(context.get_available_memory(), Some(1024 * 1024 * 1024));
1626 assert_eq!(context.get_total_memory(), Some(4 * 1024 * 1024 * 1024));
1627 }
1628
1629 #[test]
1630 fn test_gpu_context_buffer_creation() {
1631 let context = GpuContext::new(GpuBackend::Cpu).expect("Operation failed");
1632
1633 let buffer = context.create_buffer::<f32>(100);
1634 assert_eq!(buffer.len(), 100);
1635
1636 let data = vec![1.0f32; 50];
1637 let buffer_from_slice = context.create_buffer_from_slice(&data);
1638 assert_eq!(buffer_from_slice.len(), 50);
1639
1640 let result = buffer_from_slice.to_vec();
1641 assert_eq!(result, data);
1642 }
1643
1644 #[test]
1645 fn test_gpu_context_unsupported_backend() {
1646 let result = GpuContext::new(GpuBackend::Cuda);
1649 match result {
1650 Err(GpuError::BackendNotAvailable(msg)) => {
1651 assert!(
1652 msg.contains("oxicuda"),
1653 "expected oxicuda guidance, got: {msg}"
1654 );
1655 }
1656 other => {
1657 panic!("expected BackendNotAvailable for retired CUDA backend, got: {other:?}")
1658 }
1659 }
1660 }
1661
1662 #[test]
1663 fn test_gpu_compiler() {
1664 let compiler_impl = Arc::new(CpuCompiler);
1665 let compiler = GpuCompiler::new(compiler_impl);
1666
1667 let kernel = compiler
1669 .compile("dummy kernel source")
1670 .expect("Operation failed");
1671 kernel.dispatch([1, 1, 1]);
1672 }
1673
1674 #[test]
1680 fn test_gpu_compiler_typed_vector_add_computes_real_result() {
1681 let compiler = GpuCompiler::new(Arc::new(CpuCompiler));
1682 let context = GpuContext::new(GpuBackend::Cpu).expect("Operation failed");
1683
1684 let a_data = [1.0f32, 2.0, 3.0, 4.0];
1685 let b_data = [10.0f32, 20.0, 30.0, 40.0];
1686 let a = context.create_buffer_from_slice(&a_data);
1687 let b = context.create_buffer_from_slice(&b_data);
1688 let result = context.create_buffer::<f32>(a_data.len());
1689
1690 let typed_kernel = compiler
1691 .compile_kernel::<f32, f32>("vector_add")
1692 .expect("a real f32 vector_add kernel should compile on the CPU backend");
1693 typed_kernel.set_buffer("a", &a);
1694 typed_kernel.set_buffer("b", &b);
1695 typed_kernel.set_buffer("result", &result);
1696 typed_kernel.dispatch([1, 1, 1]);
1697
1698 let computed = result.to_vec();
1699 assert_eq!(
1700 computed,
1701 vec![11.0, 22.0, 33.0, 44.0],
1702 "vector_add must actually compute a + b, not leave the output untouched"
1703 );
1704 }
1705
1706 #[test]
1709 fn test_gpu_compiler_typed_unknown_kernel_returns_honest_error() {
1710 let compiler = GpuCompiler::new(Arc::new(CpuCompiler));
1711 let result = compiler.compile_kernel::<f32, f32>("totally_unknown_kernel_name");
1712 assert!(
1713 result.is_err(),
1714 "an unrecognized kernel name must return an error, not a fabricated success"
1715 );
1716 }
1717
1718 #[test]
1719 fn test_gpu_context_execute() {
1720 let context = GpuContext::new(GpuBackend::Cpu).expect("Operation failed");
1721
1722 let result = context.execute(|compiler| compiler.compile("test kernel").is_ok());
1723
1724 assert!(result);
1725 }
1726
1727 #[test]
1728 fn test_gpu_context_kernel_registry() {
1729 let context = GpuContext::new(GpuBackend::Cpu).expect("Operation failed");
1730
1731 let result = context.get_kernel("non_existent_kernel");
1733 assert!(result.is_err());
1734 match result {
1735 Err(GpuError::KernelNotFound(_)) => {}
1736 _ => panic!("Expected KernelNotFound error"),
1737 }
1738 }
1739
1740 #[test]
1741 fn test_cpu_buffer_implementation() {
1742 let buffer = CpuBuffer::new(256);
1743 assert_eq!(buffer.data.len(), 256);
1744
1745 assert!(buffer.data.iter().all(|&b| b == 0));
1747 }
1748
1749 #[test]
1750 fn test_gpuerror_display() {
1751 let error = GpuError::BackendNotAvailable("CUDA".to_string());
1752 assert_eq!(error.to_string(), "GPU backend CUDA is not available");
1753
1754 let error = GpuError::OutOfMemory("allocation failed".to_string());
1755 assert_eq!(error.to_string(), "GPU out of memory: allocation failed");
1756
1757 let error = GpuError::KernelCompilationError("syntax error".to_string());
1758 assert_eq!(error.to_string(), "Kernel compilation error: syntax error");
1759
1760 let error = GpuError::KernelNotFound("gemm".to_string());
1761 assert_eq!(error.to_string(), "Kernel not found: gemm");
1762 }
1763
1764 #[test]
1765 fn test_backend_equality() {
1766 assert_eq!(GpuBackend::Cuda, GpuBackend::Cuda);
1767 assert_ne!(GpuBackend::Cuda, GpuBackend::Rocm);
1768
1769 let backend = GpuBackend::Metal;
1771 let cloned = backend;
1772 let copied = backend;
1773 assert_eq!(backend, cloned);
1774 assert_eq!(backend, copied);
1775 }
1776
1777 #[test]
1778 fn test_backend_hash() {
1779 use std::collections::HashSet;
1780
1781 let mut set = HashSet::new();
1782 set.insert(GpuBackend::Cuda);
1783 set.insert(GpuBackend::Rocm);
1784 set.insert(GpuBackend::Cuda); assert_eq!(set.len(), 2); assert!(set.contains(&GpuBackend::Cuda));
1788 assert!(set.contains(&GpuBackend::Rocm));
1789 }
1790
1791 #[test]
1792 fn test_gpu_buffer_debug_clone() {
1793 let inner = Arc::new(CpuBuffer::new(16));
1794 let buffer = GpuBuffer::<f32>::new(inner, 4);
1795
1796 let debug_str = format!("{:?}", buffer);
1798 assert!(debug_str.contains("GpuBuffer"));
1799 assert!(debug_str.contains("size"));
1800
1801 let cloned = buffer.clone();
1803 assert_eq!(cloned.len(), buffer.len());
1804 assert_eq!(cloned.len(), 4);
1805
1806 let data = vec![1.0f32, 2.0, 3.0, 4.0];
1808 let _ = buffer.copy_from_host(&data);
1809
1810 let mut result = vec![0.0f32; 4];
1811 let _ = cloned.copy_to_host(&mut result);
1812 assert_eq!(result, data);
1813 }
1814
1815 #[test]
1816 fn test_gpu_context_debug() {
1817 let context = GpuContext::new(GpuBackend::Cpu).expect("Failed to create context");
1818
1819 let debug_str = format!("{:?}", context);
1821 assert!(debug_str.contains("GpuContext"));
1822 assert!(debug_str.contains("backend"));
1823 assert!(debug_str.contains("Cpu"));
1824 }
1825
1826 #[test]
1827 fn test_gpu_context_batch_dispatch() {
1828 let context = GpuContext::new(GpuBackend::Cpu).expect("Failed to create CPU context");
1830
1831 let begin_result = context.begin_batch();
1833 assert!(
1834 begin_result.is_ok(),
1835 "begin_batch should succeed on CPU backend"
1836 );
1837
1838 let dispatch_result = context.execute(|compiler| {
1840 compiler.compile("dummy kernel source").map(|kernel| {
1841 kernel.dispatch([4, 1, 1]);
1842 })
1843 });
1844 assert!(
1845 dispatch_result.is_ok(),
1846 "kernel dispatch inside batch should succeed"
1847 );
1848
1849 let end_result = context.end_batch();
1851 assert!(
1852 end_result.is_ok(),
1853 "end_batch should succeed on CPU backend"
1854 );
1855 }
1856
1857 #[test]
1858 fn test_gpu_context_gpu_sync() {
1859 let context = GpuContext::new(GpuBackend::Cpu).expect("Failed to create CPU context");
1861
1862 let result = context.gpu_sync();
1864 assert!(result.is_ok(), "gpu_sync should return Ok on CPU backend");
1865 }
1866
1867 #[test]
1868 fn test_gpu_kernel_dispatch_no_wait() {
1869 let kernel = Arc::new(CpuKernel);
1871 let handle = GpuKernelHandle::new(kernel);
1872
1873 let buffer = GpuBuffer::<f32>::new(Arc::new(CpuBuffer::new(16)), 4);
1875 handle.set_buffer("input", &buffer);
1876 handle.set_u32("size", 4);
1877
1878 handle.dispatch_no_wait([4, 1, 1]);
1880 }
1881
1882 #[test]
1883 fn test_gpu_device_get_info_cpu() {
1884 let device = GpuDevice::new(GpuBackend::Cpu, 0);
1885 let info = device
1886 .get_info()
1887 .expect("get_info should succeed for the CPU backend");
1888
1889 assert_eq!(info.backend, GpuBackend::Cpu);
1891 assert_eq!(info.device_name, "CPU");
1892 assert_eq!(info.device_type, "CPU");
1893
1894 assert!(!info.device_name.is_empty());
1897 assert!(!info.compute_capability.is_empty());
1898 assert!(info.max_work_group_size >= 1);
1899 assert!(info.supports_fp64);
1900 assert!(info.supports_fp16);
1901 }
1902
1903 #[test]
1904 fn test_gpu_device_info_is_deterministic_per_backend() {
1905 for backend in [
1907 GpuBackend::Cpu,
1908 GpuBackend::Cuda,
1909 GpuBackend::Rocm,
1910 GpuBackend::Wgpu,
1911 GpuBackend::Metal,
1912 GpuBackend::OpenCL,
1913 ] {
1914 let device = GpuDevice::new(backend, 0);
1915 let info = device.get_info().expect("get_info should not fail");
1916 let info_again = device.get_info().expect("get_info should not fail");
1917
1918 assert_eq!(info.backend, backend);
1919 assert!(!info.device_name.is_empty());
1920 assert!(!info.device_type.is_empty());
1921 assert!(!info.compute_capability.is_empty());
1922 assert!(info.max_work_group_size >= 1);
1923
1924 assert_eq!(info.device_name, info_again.device_name);
1926 assert_eq!(info.max_work_group_size, info_again.max_work_group_size);
1927 assert_eq!(info.supports_fp64, info_again.supports_fp64);
1928 }
1929 }
1930}