66 lines
1.8 KiB
Rust
66 lines
1.8 KiB
Rust
//! RustyBooks GPU Integration for RustyTorch
|
|
//!
|
|
//! This crate provides integration between the RustyBooks GPU acceleration crates
|
|
//! and the RustyTorch ML framework.
|
|
//!
|
|
//! # Features
|
|
//!
|
|
//! - `direct-storage` - NVMe→GPU direct transfers via gpu-direct
|
|
//! - `profiling` - Extended kernel profiling via gpu-profiler
|
|
//! - `tma` - Tensor Memory Accelerator operations via tensor-accelerator
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_rustybooks::memory::DirectStoragePool;
|
|
//! use rtx_rustybooks::profiler::ExtendedKernelProfiler;
|
|
//! use rtx_rustybooks::tensor::TmaOps;
|
|
//!
|
|
//! // Direct model loading (bypasses CPU)
|
|
//! let pool = DirectStoragePool::new()?;
|
|
//! pool.load_from_storage("model.safetensors", &mut weights).await?;
|
|
//!
|
|
//! // Profiled kernel execution
|
|
//! let profiler = ExtendedKernelProfiler::new()?;
|
|
//! let result = profiler.profile_launch("matmul", || tensor_op());
|
|
//! println!("Occupancy: {:.1}%", result.warp_occupancy() * 100.0);
|
|
//! ```
|
|
|
|
pub mod error;
|
|
|
|
#[cfg(feature = "direct-storage")]
|
|
pub mod memory;
|
|
|
|
#[cfg(feature = "profiling")]
|
|
pub mod profiler;
|
|
|
|
#[cfg(feature = "tma")]
|
|
pub mod tensor;
|
|
|
|
// Re-exports for convenience
|
|
#[cfg(feature = "direct-storage")]
|
|
pub use memory::{DirectBuffer, DirectStoragePool, DirectTransfer};
|
|
|
|
#[cfg(feature = "profiling")]
|
|
pub use profiler::{ExtendedKernelProfiler, ProfilingResult, WarpMetrics};
|
|
|
|
#[cfg(feature = "tma")]
|
|
pub use tensor::{TileConfig, TmaOps};
|
|
|
|
/// Check if RustyBooks integration is available
|
|
pub fn is_available() -> bool {
|
|
#[cfg(any(feature = "direct-storage", feature = "profiling", feature = "tma"))]
|
|
{
|
|
true
|
|
}
|
|
#[cfg(not(any(feature = "direct-storage", feature = "profiling", feature = "tma")))]
|
|
{
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Get version information
|
|
pub fn version() -> &'static str {
|
|
env!("CARGO_PKG_VERSION")
|
|
}
|