102 lines
2.7 KiB
Rust
102 lines
2.7 KiB
Rust
//! # `RustyTorch`++ `CubeCL` Integration
|
|
//!
|
|
//! This crate provides GPU kernel compilation using `CubeCL`, enabling:
|
|
//! - Writing GPU kernels in Rust syntax
|
|
//! - Single-source compilation to CUDA/WebGPU/ROCm/Metal/Vulkan
|
|
//! - Type-safe kernel parameters
|
|
//! - Automatic vectorization and optimization
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! `CubeCL` kernels are written using the `#[cube]` proc macro:
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use cubecl::prelude::*;
|
|
//!
|
|
//! #[cube(launch)]
|
|
//! fn vector_add<F: Float>(a: &Tensor<F>, b: &Tensor<F>, out: &mut Tensor<F>) {
|
|
//! let idx = ABSOLUTE_POS;
|
|
//! if idx < out.len() {
|
|
//! out[idx] = a[idx] + b[idx];
|
|
//! }
|
|
//! }
|
|
//! ```
|
|
//!
|
|
//! ## Backend Integration
|
|
//!
|
|
//! The `CubeclBackend` type implements the `Backend` trait from `rtx-backend`:
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_cubecl::CubeclBackend;
|
|
//! use rtx_backend::Backend;
|
|
//!
|
|
//! // Use CubeCL as compute backend
|
|
//! type MyBackend = CubeclBackend;
|
|
//!
|
|
//! let device = CubeclDevice::default();
|
|
//! let tensor = MyBackend::zeros::<2>([3, 4], &device);
|
|
//! ```
|
|
//!
|
|
//! ## Backends
|
|
//!
|
|
//! Enable backends via feature flags:
|
|
//! - `wgpu` - WebGPU (cross-platform, browsers, desktop)
|
|
//! - `cuda` - NVIDIA CUDA
|
|
//! - `hip` - AMD ROCm/HIP
|
|
//! - `vulkan` - Vulkan via WebGPU
|
|
//! - `metal` - Apple Metal via WebGPU
|
|
//!
|
|
//! ## Tiered Kernel Strategy
|
|
//!
|
|
//! `RustyTorch`++ maintains both `CubeCL` and hand-crafted kernels:
|
|
//!
|
|
//! - **`CubeCL` (Tier 1)**: Portable, type-safe, single-source (~60% of ops)
|
|
//! - Elementwise: add, sub, mul, div, exp, log
|
|
//! - Activations: relu, gelu, sigmoid, tanh
|
|
//! - Reductions: sum, mean, max, min
|
|
//!
|
|
//! - **Hand-crafted (Tier 2)**: Maximum performance for LLM-critical ops (~40%)
|
|
//! - `FlashAttention`
|
|
//! - Large matrix multiplications (cuBLAS)
|
|
//! - Convolutions (cuDNN)
|
|
//!
|
|
//! The `rtx-kernel-bench` crate provides A/B benchmarking between them.
|
|
|
|
#![allow(clippy::module_name_repetitions)]
|
|
#![allow(clippy::must_use_candidate)]
|
|
#![allow(clippy::missing_errors_doc)]
|
|
|
|
pub mod backend;
|
|
pub mod client;
|
|
pub mod device;
|
|
pub mod error;
|
|
pub mod kernels;
|
|
pub mod ops;
|
|
pub mod runtime;
|
|
pub mod tensor;
|
|
|
|
// Re-export CubeCL types for convenience
|
|
pub use cubecl::prelude::*;
|
|
pub use cubecl::{CubeCount, CubeDim, Runtime};
|
|
|
|
// Re-export our types
|
|
pub use backend::CubeclBackend;
|
|
pub use client::{CubeclClient, TensorHandle};
|
|
pub use device::CubeclDevice;
|
|
pub use error::{CubeclError, Result};
|
|
pub use runtime::{CubeclRuntime, RuntimeBackend};
|
|
pub use tensor::{CubeclTensor, CubeclTensorPrimitive};
|
|
|
|
/// Version information
|
|
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_version() {
|
|
assert!(!VERSION.is_empty());
|
|
}
|
|
}
|