morok-device 0.1.0-alpha.2

Device abstraction layer for the Morok ML compiler
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

//! Device abstraction layer for tensor operations.
//!
//! This module provides a clean abstraction over different compute devices (CPU, CUDA, etc.)
//! with support for:
//! - Lazy buffer allocation
//! - Buffer views with zero-copy slicing
//! - LRU caching of allocations for performance
//! - Device-agnostic copy operations
//!
//! # Examples
//!
//! ```no_run
//! use morok_device::allocator::BufferOptions;
//! use morok_device::{Buffer, registry};
//! use morok_dtype::DType;
//!
//! // Get a CPU device
//! let cpu = registry::cpu().unwrap();
//!
//! // Create a buffer with lazy allocation
//! let buffer = Buffer::new(cpu, DType::Float32, vec![10, 10], BufferOptions::default());
//!
//! // Allocation happens on first use
//! buffer.ensure_allocated().unwrap();
//! ```

pub mod allocator;
pub mod buffer;
pub mod device;
pub mod error;
pub mod queue;
pub mod registry;
pub mod sync;

pub use buffer::{Buffer, BufferId};
pub use error::{Error, Result};
pub use queue::{DynQueue, ExecParams, HardwareQueue, Program, QueueFactory};
pub use sync::{CpuTimelineSignal, TimelineSignal};

#[cfg(test)]
mod test;

// Re-export commonly used types
#[cfg(feature = "cuda")]
pub use allocator::CudaAllocator;
pub use allocator::{Allocator, BufferOptions, CpuAllocator};
#[cfg(feature = "cuda")]
pub use registry::cuda;
pub use registry::{DeviceSpec, cpu, get_device};