ferrotensor_backend/lib.rs
1//! wgpu compute backend for Ferrotensor: GPU tensors and hand-written WGSL kernels.
2//!
3//! [`WgpuContext`] opens a GPU (Metal, Vulkan, DX12), [`UnifiedTensor`] holds
4//! data on it and [`ComputeEngine`] compiles the kernels and dispatches them:
5//! matmul (plain and transposed), element-wise ops, softmax, RMSNorm, RoPE,
6//! fused flash attention (forward and backward, with GQA, masks, dropout,
7//! KV-cache offsets and sliding windows), losses, optimisers and dequantisation.
8//!
9//! ```no_run
10//! use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
11//!
12//! let ctx = pollster::block_on(WgpuContext::new());
13//! let engine = ComputeEngine::new(&ctx);
14//! let a = UnifiedTensor::new(&ctx, &[1.0, 2.0, 3.0, 4.0], &[2, 2]);
15//! let b = UnifiedTensor::new(&ctx, &[1.0, 0.0, 0.0, 1.0], &[2, 2]);
16//! let c = UnifiedTensor::empty(&ctx, &[2, 2]);
17//! engine.matmul(&ctx, &a, &b, &c);
18//! assert_eq!(pollster::block_on(c.to_vec(&ctx)), vec![1.0, 2.0, 3.0, 4.0]);
19//! ```
20//!
21//! Kernels are checked against CPU references in `tests/gpu_kernels.rs`
22//! (run with `cargo test -p ferrotensor-backend --test gpu_kernels -- --ignored
23//! --test-threads=1`).
24
25/// The `wgpu` version this crate is built on, re-exported for callers.
26pub use wgpu; // Re-export wgpu for downstream crates
27
28mod context;
29mod engine;
30/// Size-keyed reuse of GPU buffers.
31pub mod pool;
32mod tensor;
33
34pub use context::WgpuContext;
35pub use engine::{AdamParams, AttentionOptions, ComputeEngine};
36pub use tensor::{DType, UnifiedTensor};