Skip to main content

oxicuda_dnn/
lib.rs

1//! # OxiCUDA DNN -- GPU-Accelerated Deep Learning Primitives
2//!
3//! This crate provides GPU-accelerated deep learning primitives,
4//! serving as a pure Rust equivalent to cuDNN.
5//!
6//! ## Modules
7//!
8//! | Module | Description |
9//! |--------|-------------|
10//! | [`error`] | Error types and `DnnResult<T>` alias |
11//! | [`types`] | Tensor descriptors, layouts, activations, conv descriptors |
12//! | [`handle`] | `DnnHandle` -- central entry point for all operations |
13//! | [`conv`] | Convolution forward / backward / fused operations |
14
15#![warn(clippy::all)]
16#![warn(missing_docs)]
17
18pub mod activation;
19pub mod attn;
20pub mod conv;
21pub mod dynamic_batch;
22pub mod error;
23pub mod handle;
24pub mod layers;
25pub mod linear;
26pub mod moe;
27pub mod norm;
28pub mod pool;
29pub mod position;
30pub mod quantize;
31pub mod resize;
32pub mod rnn;
33pub mod types;
34
35pub(crate) mod ptx_helpers;
36pub(crate) mod tensor_util;
37
38// ---------------------------------------------------------------------------
39// Re-exports
40// ---------------------------------------------------------------------------
41
42pub use activation::{SwiGlu, SwiGluConfig};
43pub use dynamic_batch::{
44    BatchConfig, BatchDecision, BatchMetrics, BatchSlot, ContinuousBatcher, DraftedToken,
45    InferenceRequest, LcgRng, PagedKvManager, PreemptionPolicy, Priority, RequestId,
46    SchedulingPolicy, SpeculativeDecoder, SpeculativeResult, TokenBudgetAllocator,
47};
48pub use error::{DnnError, DnnResult};
49pub use handle::DnnHandle;
50pub use position::{AlibiBias, DnnRng, Rope, RopeConfig, alibi_slope};
51pub use types::{
52    Activation, ConvAlgorithm, ConvolutionDescriptor, TensorDesc, TensorDescMut, TensorLayout,
53    pool_output_size,
54};
55
56/// Prelude module for convenient glob imports.
57///
58/// ```rust,no_run
59/// use oxicuda_dnn::prelude::*;
60/// ```
61pub mod prelude {
62    pub use crate::error::{DnnError, DnnResult};
63    pub use crate::handle::DnnHandle;
64    pub use crate::types::{
65        Activation, ConvAlgorithm, ConvolutionDescriptor, TensorDesc, TensorDescMut, TensorLayout,
66        pool_output_size,
67    };
68}