tenflowers_core/device/
context.rs

1use crate::{Device, Result, TensorError};
2use std::collections::HashMap;
3use std::sync::Arc;
4
5/// Device-specific memory allocator trait
6pub trait DeviceAllocator: Send + Sync {
7    /// Allocate memory on the device
8    fn allocate(&self, size: usize) -> Result<*mut u8>;
9
10    /// Deallocate memory on the device
11    ///
12    /// # Safety
13    /// The caller must ensure that:
14    /// - The pointer was allocated by this allocator
15    /// - The size matches the original allocation
16    /// - The pointer is not used after deallocation
17    unsafe fn deallocate(&self, ptr: *mut u8, size: usize);
18
19    /// Copy memory from host to device
20    fn copy_from_host(&self, dst: *mut u8, src: &[u8]) -> Result<()>;
21
22    /// Copy memory from device to host
23    fn copy_to_host(&self, dst: &mut [u8], src: *const u8) -> Result<()>;
24
25    /// Copy memory within device
26    fn copy_device_to_device(&self, dst: *mut u8, src: *const u8, size: usize) -> Result<()>;
27}
28
29/// Device-specific compute context
30pub trait DeviceContext: Send + Sync {
31    /// Get the device
32    fn device(&self) -> &Device;
33
34    /// Get the allocator for this device
35    fn allocator(&self) -> &dyn DeviceAllocator;
36
37    /// Synchronize device operations
38    fn synchronize(&self) -> Result<()>;
39
40    /// Create a compute stream/queue
41    fn create_stream(&self) -> Result<Box<dyn DeviceStream>>;
42
43    /// Get device properties
44    fn properties(&self) -> DeviceProperties;
45}
46
47/// Device stream for async operations
48pub trait DeviceStream: Send + Sync {
49    /// Enqueue a kernel launch
50    fn launch_kernel(&self, kernel: &dyn DeviceKernel, args: KernelArgs) -> Result<()>;
51
52    /// Synchronize the stream
53    fn synchronize(&self) -> Result<()>;
54
55    /// Check if operations are complete
56    fn is_complete(&self) -> bool;
57}
58
59/// Device kernel interface
60pub trait DeviceKernel: Send + Sync {
61    /// Get kernel name
62    fn name(&self) -> &str;
63
64    /// Get required shared memory size
65    fn shared_memory_size(&self) -> usize;
66
67    /// Get optimal block size
68    fn optimal_block_size(&self) -> usize;
69}
70
71/// Kernel launch arguments
72pub struct KernelArgs {
73    pub grid_size: [u32; 3],
74    pub block_size: [u32; 3],
75    pub shared_memory: usize,
76    pub inputs: Vec<*const u8>,
77    pub outputs: Vec<*mut u8>,
78    pub params: HashMap<String, KernelParam>,
79}
80
81/// Kernel parameter types
82#[derive(Clone)]
83pub enum KernelParam {
84    I32(i32),
85    U32(u32),
86    F32(f32),
87    F64(f64),
88    Ptr(*const u8),
89}
90
91/// Device properties
92#[derive(Debug, Clone)]
93pub struct DeviceProperties {
94    pub name: String,
95    pub total_memory: usize,
96    pub compute_capability: (u32, u32),
97    pub max_threads_per_block: u32,
98    pub max_blocks: [u32; 3],
99    pub shared_memory_per_block: usize,
100    pub warp_size: u32,
101}
102
103/// CPU device context implementation
104pub struct CpuContext {
105    allocator: CpuAllocator,
106}
107
108impl CpuContext {
109    pub fn new() -> Self {
110        Self {
111            allocator: CpuAllocator,
112        }
113    }
114}
115
116impl Default for CpuContext {
117    fn default() -> Self {
118        Self::new()
119    }
120}
121
122impl DeviceContext for CpuContext {
123    fn device(&self) -> &Device {
124        &Device::Cpu
125    }
126
127    fn allocator(&self) -> &dyn DeviceAllocator {
128        &self.allocator
129    }
130
131    fn synchronize(&self) -> Result<()> {
132        // CPU operations are synchronous
133        Ok(())
134    }
135
136    fn create_stream(&self) -> Result<Box<dyn DeviceStream>> {
137        Ok(Box::new(CpuStream))
138    }
139
140    fn properties(&self) -> DeviceProperties {
141        DeviceProperties {
142            name: "CPU".to_string(),
143            total_memory: sys_info::mem_info()
144                .map(|info| info.total as usize * 1024)
145                .unwrap_or(0),
146            compute_capability: (1, 0),
147            max_threads_per_block: 1,
148            max_blocks: [1, 1, 1],
149            shared_memory_per_block: 0,
150            warp_size: 1,
151        }
152    }
153}
154
155/// CPU allocator
156struct CpuAllocator;
157
158impl DeviceAllocator for CpuAllocator {
159    fn allocate(&self, size: usize) -> Result<*mut u8> {
160        if size == 0 {
161            return Ok(std::ptr::null_mut());
162        }
163
164        let layout = std::alloc::Layout::from_size_align(size, 64)
165            .map_err(|e| TensorError::allocation_error_simple(e.to_string()))?;
166
167        let ptr = unsafe { std::alloc::alloc(layout) };
168        if ptr.is_null() {
169            Err(TensorError::allocation_error_simple(format!(
170                "Failed to allocate {size} bytes"
171            )))
172        } else {
173            Ok(ptr)
174        }
175    }
176
177    unsafe fn deallocate(&self, ptr: *mut u8, size: usize) {
178        if !ptr.is_null() && size > 0 {
179            let layout = std::alloc::Layout::from_size_align_unchecked(size, 64);
180            std::alloc::dealloc(ptr, layout);
181        }
182    }
183
184    fn copy_from_host(&self, dst: *mut u8, src: &[u8]) -> Result<()> {
185        unsafe {
186            std::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len());
187        }
188        Ok(())
189    }
190
191    fn copy_to_host(&self, dst: &mut [u8], src: *const u8) -> Result<()> {
192        unsafe {
193            std::ptr::copy_nonoverlapping(src, dst.as_mut_ptr(), dst.len());
194        }
195        Ok(())
196    }
197
198    fn copy_device_to_device(&self, dst: *mut u8, src: *const u8, size: usize) -> Result<()> {
199        unsafe {
200            std::ptr::copy_nonoverlapping(src, dst, size);
201        }
202        Ok(())
203    }
204}
205
206/// CPU stream (synchronous)
207struct CpuStream;
208
209impl DeviceStream for CpuStream {
210    fn launch_kernel(&self, _kernel: &dyn DeviceKernel, _args: KernelArgs) -> Result<()> {
211        // CPU kernels are executed directly, not through stream
212        Ok(())
213    }
214
215    fn synchronize(&self) -> Result<()> {
216        Ok(())
217    }
218
219    fn is_complete(&self) -> bool {
220        true
221    }
222}
223
224/// Device manager for context creation and caching
225pub struct DeviceManager {
226    contexts: std::sync::RwLock<HashMap<Device, Arc<dyn DeviceContext>>>,
227}
228
229impl DeviceManager {
230    pub fn new() -> Self {
231        Self {
232            contexts: std::sync::RwLock::new(HashMap::new()),
233        }
234    }
235}
236
237impl Default for DeviceManager {
238    fn default() -> Self {
239        Self::new()
240    }
241}
242
243impl DeviceManager {
244    /// Get or create a device context
245    pub fn get_context(&self, device: &Device) -> Result<Arc<dyn DeviceContext>> {
246        // Check cache first
247        {
248            let contexts = self.contexts.read().map_err(|_| {
249                TensorError::invalid_operation_simple("device contexts lock poisoned".to_string())
250            })?;
251            if let Some(ctx) = contexts.get(device) {
252                return Ok(Arc::clone(ctx));
253            }
254        }
255
256        // Create new context
257        let context: Arc<dyn DeviceContext> = match device {
258            Device::Cpu => Arc::new(CpuContext::new()),
259            #[cfg(feature = "gpu")]
260            Device::Gpu(id) => Arc::new(GpuContext::new(*id)?),
261            #[cfg(feature = "rocm")]
262            Device::Rocm(id) => Arc::new(GpuContext::new(*id)?), // Use same GPU context for ROCm
263        };
264
265        // Cache it
266        {
267            let mut contexts = self.contexts.write().map_err(|_| {
268                TensorError::invalid_operation_simple("device contexts lock poisoned".to_string())
269            })?;
270            contexts.insert(*device, Arc::clone(&context));
271        }
272
273        Ok(context)
274    }
275}
276
277// Global device manager
278lazy_static::lazy_static! {
279    pub static ref DEVICE_MANAGER: DeviceManager = DeviceManager::new();
280}
281
282/// GPU device context implementation
283#[cfg(feature = "gpu")]
284pub struct GpuContext {
285    device: Arc<wgpu::Device>,
286    queue: Arc<wgpu::Queue>,
287    adapter: wgpu::Adapter,
288    allocator: GpuAllocator,
289    device_enum: Device,
290}
291
292#[cfg(feature = "gpu")]
293impl GpuContext {
294    pub fn new(device_id: usize) -> Result<Self> {
295        pollster::block_on(Self::new_async(device_id))
296    }
297
298    async fn new_async(device_id: usize) -> Result<Self> {
299        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
300
301        let adapter = instance
302            .request_adapter(&wgpu::RequestAdapterOptions {
303                power_preference: wgpu::PowerPreference::HighPerformance,
304                compatible_surface: None,
305                force_fallback_adapter: false,
306                apply_limit_buckets: false,
307            })
308            .await
309            .map_err(|_e| {
310                TensorError::device_error_simple("Failed to find suitable GPU adapter".to_string())
311            })?;
312
313        let (device, queue) = adapter
314            .request_device(&wgpu::DeviceDescriptor {
315                label: Some(&format!("TenfloweRS GPU Device {device_id}")),
316                required_features: wgpu::Features::empty(),
317                required_limits: wgpu::Limits::default(),
318                memory_hints: wgpu::MemoryHints::Performance,
319                experimental_features: wgpu::ExperimentalFeatures::default(),
320                trace: wgpu::Trace::default(),
321            })
322            .await
323            .map_err(|e| {
324                TensorError::device_error_simple(format!("Failed to create GPU device: {e}"))
325            })?;
326
327        let device = Arc::new(device);
328        let queue = Arc::new(queue);
329        let allocator = GpuAllocator::new(Arc::clone(&device), Arc::clone(&queue));
330
331        Ok(Self {
332            device,
333            queue,
334            adapter,
335            allocator,
336            device_enum: Device::Gpu(device_id),
337        })
338    }
339
340    /// Real capability data read directly from the live `wgpu::Adapter`.
341    ///
342    /// This is a thin, honest snapshot of exactly what `wgpu::Adapter` exposes: the raw
343    /// `AdapterInfo` (name/vendor id/backend/driver strings), `Limits` (workgroup size,
344    /// shared-memory size, buffer size caps, ...), and `Features` (optional capability
345    /// flags such as `SHADER_F16`). No reinterpretation or vendor-specific guessing
346    /// happens here — that mapping lives in callers (e.g.
347    /// `crate::gpu::advanced_kernel_manager`) that need it.
348    pub fn adapter_capabilities(&self) -> GpuAdapterCapabilities {
349        GpuAdapterCapabilities {
350            info: self.adapter.get_info(),
351            limits: self.adapter.limits(),
352            features: self.adapter.features(),
353        }
354    }
355}
356
357#[cfg(feature = "gpu")]
358impl DeviceContext for GpuContext {
359    fn device(&self) -> &Device {
360        &self.device_enum
361    }
362
363    fn allocator(&self) -> &dyn DeviceAllocator {
364        &self.allocator
365    }
366
367    fn synchronize(&self) -> Result<()> {
368        self.device.poll(wgpu::PollType::wait_indefinitely()).ok();
369        Ok(())
370    }
371
372    fn create_stream(&self) -> Result<Box<dyn DeviceStream>> {
373        Ok(Box::new(GpuStream::new(
374            Arc::clone(&self.device),
375            Arc::clone(&self.queue),
376        )))
377    }
378
379    fn properties(&self) -> DeviceProperties {
380        let info = self.adapter.get_info();
381        let limits = self.adapter.limits();
382
383        DeviceProperties {
384            name: info.name,
385            total_memory: 0,            // WGPU doesn't expose memory info
386            compute_capability: (1, 0), // Generic capability
387            max_threads_per_block: limits.max_compute_workgroup_size_x,
388            max_blocks: [
389                limits.max_compute_workgroups_per_dimension,
390                limits.max_compute_workgroups_per_dimension,
391                limits.max_compute_workgroups_per_dimension,
392            ],
393            shared_memory_per_block: limits.max_compute_workgroup_storage_size as usize,
394            warp_size: 32, // Common warp size
395        }
396    }
397}
398
399/// GPU allocator implementation
400#[cfg(feature = "gpu")]
401struct GpuAllocator {
402    device: Arc<wgpu::Device>,
403    queue: Arc<wgpu::Queue>,
404}
405
406#[cfg(feature = "gpu")]
407impl GpuAllocator {
408    fn new(device: Arc<wgpu::Device>, queue: Arc<wgpu::Queue>) -> Self {
409        Self { device, queue }
410    }
411}
412
413#[cfg(feature = "gpu")]
414impl DeviceAllocator for GpuAllocator {
415    fn allocate(&self, _size: usize) -> Result<*mut u8> {
416        // WGPU does not expose raw host-side pointers for GPU memory.
417        // GPU buffers must be managed through wgpu::Buffer / wgpu::Queue::write_buffer.
418        // Use the wgpu-native buffer API (GpuContextInfo / wgpu::Device::create_buffer)
419        // for GPU memory management instead of this raw-pointer interface.
420        Err(TensorError::unsupported_operation_simple(
421            "raw-pointer GPU allocation is not supported; use wgpu::Buffer API for GPU memory management".to_string(),
422        ))
423    }
424
425    unsafe fn deallocate(&self, _ptr: *mut u8, _size: usize) {
426        // No allocation was ever performed through this interface, so there is
427        // nothing to free.  wgpu::Buffer handles its own lifetime via Drop.
428    }
429
430    fn copy_from_host(&self, _dst: *mut u8, _src: &[u8]) -> Result<()> {
431        Err(TensorError::unsupported_operation_simple(
432            "raw-pointer GPU copy_from_host is not supported; use wgpu::Queue::write_buffer for host-to-GPU transfers".to_string(),
433        ))
434    }
435
436    fn copy_to_host(&self, _dst: &mut [u8], _src: *const u8) -> Result<()> {
437        Err(TensorError::unsupported_operation_simple(
438            "raw-pointer GPU copy_to_host is not supported; use wgpu::Buffer::map_read for GPU-to-host transfers".to_string(),
439        ))
440    }
441
442    fn copy_device_to_device(&self, _dst: *mut u8, _src: *const u8, _size: usize) -> Result<()> {
443        Err(TensorError::unsupported_operation_simple(
444            "raw-pointer GPU copy_device_to_device is not supported; use wgpu::CommandEncoder::copy_buffer_to_buffer for GPU-to-GPU copies".to_string(),
445        ))
446    }
447}
448
449/// GPU stream implementation
450#[cfg(feature = "gpu")]
451struct GpuStream {
452    device: Arc<wgpu::Device>,
453    queue: Arc<wgpu::Queue>,
454}
455
456#[cfg(feature = "gpu")]
457impl GpuStream {
458    fn new(device: Arc<wgpu::Device>, queue: Arc<wgpu::Queue>) -> Self {
459        Self { device, queue }
460    }
461}
462
463#[cfg(feature = "gpu")]
464impl DeviceStream for GpuStream {
465    fn launch_kernel(&self, _kernel: &dyn DeviceKernel, _args: KernelArgs) -> Result<()> {
466        // Kernel execution will be implemented with compute shaders
467        Ok(())
468    }
469
470    fn synchronize(&self) -> Result<()> {
471        self.device.poll(wgpu::PollType::wait_indefinitely()).ok();
472        Ok(())
473    }
474
475    fn is_complete(&self) -> bool {
476        // For WGPU, operations are generally synchronous
477        true
478    }
479}
480
481/// Helper function to get GPU context with backend selection
482#[cfg(feature = "gpu")]
483pub fn get_gpu_context(device_id: usize) -> Result<GpuContextInfo> {
484    let device = Device::Gpu(device_id);
485    let ctx = DEVICE_MANAGER.get_context(&device)?;
486
487    // Create a new GPU context if needed (simple approach)
488    let gpu_ctx = GpuContext::new(device_id)?;
489    Ok(GpuContextInfo {
490        device: gpu_ctx.device.clone(),
491        queue: gpu_ctx.queue.clone(),
492    })
493}
494
495/// Look up the real `wgpu::Adapter` capabilities for a GPU device by id.
496///
497/// Mirrors [`get_gpu_context`]'s "create on demand" approach: constructs a `GpuContext`
498/// for `device_id` (which internally calls `wgpu::Instance::request_adapter`) and returns
499/// its adapter's `AdapterInfo`/`Limits`/`Features` verbatim, with no reinterpretation.
500/// Callers that need vendor/feature data beyond what [`DeviceProperties`] captures (e.g.
501/// [`crate::gpu::advanced_kernel_manager`]) should use this instead of re-deriving their
502/// own adapter handle.
503#[cfg(feature = "gpu")]
504pub fn get_gpu_adapter_capabilities(device_id: usize) -> Result<GpuAdapterCapabilities> {
505    let gpu_ctx = GpuContext::new(device_id)?;
506    Ok(gpu_ctx.adapter_capabilities())
507}
508
509/// Enhanced GPU context with backend selection capabilities
510#[cfg(any(feature = "gpu", feature = "cudnn"))]
511pub fn get_enhanced_gpu_context(device_id: usize) -> Result<EnhancedGpuContext> {
512    // Check if cuDNN is available and requested
513    #[cfg(all(feature = "cudnn", any(target_os = "linux", target_os = "windows")))]
514    {
515        if crate::gpu::cudnn::CudnnContext::is_available() {
516            let mut cudnn_ctx = crate::gpu::cudnn::global_cudnn_context();
517            let cudnn_handle = cudnn_ctx.get_handle(device_id)?;
518
519            // Also get WGPU context for fallback
520            let wgpu_ctx = get_gpu_context(device_id)?;
521
522            return Ok(EnhancedGpuContext {
523                device_id,
524                backend: GpuBackend::CuDNN(cudnn_handle),
525                wgpu_fallback: Some(wgpu_ctx),
526            });
527        }
528    }
529
530    // Fall back to WGPU
531    #[cfg(feature = "gpu")]
532    {
533        let wgpu_ctx = get_gpu_context(device_id)?;
534        Ok(EnhancedGpuContext {
535            device_id,
536            backend: GpuBackend::WGPU(wgpu_ctx),
537            wgpu_fallback: None,
538        })
539    }
540    #[cfg(not(feature = "gpu"))]
541    {
542        Err(TensorError::unsupported_operation_simple(
543            "No GPU backend available",
544        ))
545    }
546}
547
548#[cfg(feature = "gpu")]
549#[derive(Debug, Clone)]
550pub struct GpuContextInfo {
551    pub device: Arc<wgpu::Device>,
552    pub queue: Arc<wgpu::Queue>,
553}
554
555/// Real, unprocessed `wgpu::Adapter` capability data.
556///
557/// Every field here comes directly from a live `wgpu::Adapter` query
558/// (`get_info()`/`limits()`/`features()`) with no fabrication or vendor-specific
559/// guessing applied. See [`GpuContext::adapter_capabilities`] and
560/// [`get_gpu_adapter_capabilities`].
561#[cfg(feature = "gpu")]
562#[derive(Debug, Clone)]
563pub struct GpuAdapterCapabilities {
564    /// Adapter identity: name, PCI vendor/device id, device type, backend, driver strings.
565    pub info: wgpu::AdapterInfo,
566    /// Real hardware/backend limits (workgroup size, shared memory, buffer sizes, ...).
567    pub limits: wgpu::Limits,
568    /// Optional capability flags supported by this adapter (e.g. `SHADER_F16`).
569    pub features: wgpu::Features,
570}
571
572/// GPU backend selection enum
573#[cfg(any(feature = "gpu", feature = "cudnn"))]
574#[derive(Debug, Clone)]
575pub enum GpuBackend {
576    #[cfg(feature = "gpu")]
577    WGPU(GpuContextInfo),
578    #[cfg(all(feature = "cudnn", any(target_os = "linux", target_os = "windows")))]
579    CuDNN(Arc<crate::gpu::cudnn::CudnnHandle>),
580}
581
582/// Enhanced GPU context with backend selection capabilities
583#[cfg(any(feature = "gpu", feature = "cudnn"))]
584#[derive(Debug, Clone)]
585pub struct EnhancedGpuContext {
586    pub device_id: usize,
587    pub backend: GpuBackend,
588    pub wgpu_fallback: Option<GpuContextInfo>,
589}
590
591#[cfg(any(feature = "gpu", feature = "cudnn"))]
592impl EnhancedGpuContext {
593    /// Check if cuDNN backend is being used
594    #[cfg(all(feature = "cudnn", any(target_os = "linux", target_os = "windows")))]
595    pub fn is_cudnn(&self) -> bool {
596        matches!(self.backend, GpuBackend::CuDNN(_))
597    }
598
599    /// Check if WGPU backend is being used
600    #[cfg(feature = "gpu")]
601    pub fn is_wgpu(&self) -> bool {
602        matches!(self.backend, GpuBackend::WGPU(_))
603    }
604
605    /// Get cuDNN handle if available
606    #[cfg(all(feature = "cudnn", any(target_os = "linux", target_os = "windows")))]
607    pub fn get_cudnn_handle(&self) -> Option<&Arc<crate::gpu::cudnn::CudnnHandle>> {
608        match &self.backend {
609            GpuBackend::CuDNN(handle) => Some(handle),
610            _ => None,
611        }
612    }
613
614    /// Get WGPU context if available
615    #[cfg(feature = "gpu")]
616    pub fn get_wgpu_context(&self) -> Option<&GpuContextInfo> {
617        match &self.backend {
618            GpuBackend::WGPU(ctx) => Some(ctx),
619            #[cfg(all(feature = "cudnn", any(target_os = "linux", target_os = "windows")))]
620            _ => self.wgpu_fallback.as_ref(),
621        }
622    }
623
624    /// Get device ID
625    pub fn device_id(&self) -> usize {
626        self.device_id
627    }
628}
629
630// Add sys-info dependency detection
631#[cfg(not(target_arch = "wasm32"))]
632extern crate sys_info;
633
634#[cfg(target_arch = "wasm32")]
635mod sys_info {
636    pub struct MemInfo {
637        pub total: u64,
638    }
639
640    pub fn mem_info() -> Result<MemInfo, &'static str> {
641        Ok(MemInfo {
642            total: 1024 * 1024 * 1024,
643        }) // 1GB default for WASM
644    }
645}
tenflowers_core/device/context.rs

tenflowers_core/device/
context.rs
