GPU Acceleration

use rumus::tensor::type">Tensor;
use rumus::nn;

"token-comment">// Move a tensor to the GPU
let x = type">Tensor::randn(&[32, 784]);
let x_gpu = x.to_gpu();

"token-comment">// Move an entire model to the GPU
model.to_gpu();

"token-comment">// All operations automatically dispatch to GPU
let output = model.forward(&x_gpu);

"token-comment">// Tensors: .to_gpu() returns a GPU-backed tensor
let x = type">Tensor::randn(&[64, 3, 28, 28]);
let x_gpu = x.to_gpu();

"token-comment">// Models: .to_gpu() moves all parameters in-place
"token-comment">// via the ModuleToGpu trait
model.to_gpu();

"token-comment">// Operations check is_gpu() at dispatch time:
"token-comment">// → GPU tensor: runs WGSL compute shader
"token-comment">// → CPU tensor: runs CPU fallback
let output = model.forward(&x_gpu);

"token-comment">// GpuContext is a OnceLock singleton — initialized
"token-comment">// lazily on first GPU operation. Never panics on
"token-comment">// missing hardware; falls back to CPU gracefully.
"token-comment">//
"token-comment">// Internally holds:
"token-comment">//   - wgpu::Device
"token-comment">//   - wgpu::Queue
"token-comment">//   - PipelineCache
"token-comment">//   - BufferPool

"token-comment">// PipelineCache: 50+ compute pipelines across 24
"token-comment">// shader modules — all validated at compile time.
"token-comment">//
"token-comment">// Pipelines are created once and cached. Each GPU op
"token-comment">// looks up its pipeline by enum variant, guaranteeing
"token-comment">// no runtime pipeline compilation stalls.
"token-comment">//
"token-comment">// Bind group layouts are shared across compatible
"token-comment">// pipelines, minimizing GPU resource usage.

"token-comment">// BufferPool: thread-safe GPU buffer cache with
"token-comment">// power-of-2 bucketing. When a buffer is dropped,
"token-comment">// it is returned to the pool for reuse.

"token-comment">// Allocation: finds the smallest power-of-2 bucket
"token-comment">// that fits the request, returns a cached buffer
"token-comment">// or allocates a new one.
let buffer = pool.allocate(1024);  "token-comment">// → 1024-byte buffer

"token-comment">// Drop: buffer is returned to its bucket
"token-comment">// automatically via the Drop trait.
drop(buffer);  "token-comment">// → recycled, not deallocated

"token-comment">// This eliminates repeated GPU memory allocation
"token-comment">// during training loops.

"token-comment">// Per-resource fences: each GPU buffer carries an
"token-comment">// AtomicUsize sentinel tracking its last submission.
"token-comment">//
"token-comment">// Before reading a buffer, RUMUS checks its fence:
"token-comment">//   - If the fence matches the current submission
"token-comment">//     index, the buffer is ready.
"token-comment">//   - Otherwise, only that buffer's fence is awaited.
"token-comment">//
"token-comment">// This avoids global pipeline stalls — independent
"token-comment">// operations on different buffers proceed in parallel.

"token-comment">// JIT Kernel Fusion: captures element-wise ops and
"token-comment">// fuses them into a single WGSL kernel at runtime.
"token-comment">//
"token-comment">// jit::compile(|| { ... }) traces the closure, builds
"token-comment">// a FusionBlock IR, and compiles a fused WGSL kernel.
"token-comment">// JitCache stores compiled pipelines for O(1) reuse.
"token-comment">//
"token-comment">// This eliminates VRAM bandwidth bottleneck by avoiding
"token-comment">// intermediate buffer round-trips between kernels.
"token-comment">//
"token-comment">// Enable with: --features jit

use rumus::jit;

let fused = jit::compile(|| {
    let y = x.mul(&w).add(&b);
    y.relu()
});

"token-comment">// Multi-GPU support: scatter-gather training and
"token-comment">// parameter sharding across multiple GPUs.
"token-comment">//
"token-comment">// MultiGpuContext::get() enumerates all available GPUs.
"token-comment">// DataParallel splits batches across devices.
"token-comment">// FSDP shards parameters for memory efficiency.
"token-comment">// AllReduceSync averages gradients across devices.
"token-comment">//
"token-comment">// Enable with: --features multi_gpu

use rumus::multi_gpu::{MultiGpuContext, DataParallel, FSDP};

let ctx = MultiGpuContext::get();
let parallel_model = DataParallel::new(&model, &ctx);
"token-comment">// or for large models:
let sharded_model = FSDP::new(&model, &ctx);

"token-comment">// Custom Ops API: register user-defined WGSL kernels.
"token-comment">//
"token-comment">// Implement the CustomOp trait and register your kernel.
"token-comment">// CustomOpCache handles dynamic bind group layout
"token-comment">// generation and pipeline caching.
"token-comment">//
"token-comment">// Enable with: --features gpu

use rumus::gpu::{CustomOp, CustomOpCache};

struct MyKernel;
impl CustomOp for MyKernel {
    fn wgsl_source(&self) -> &str { include_str!("my_kernel.wgsl") }
    fn dispatch(&self, encoder: &mut Encoder, buffers: &[&Buffer]) { "token-comment">/* ... */ }
}

use rumus::tensor::type">Tensor;
use rumus::nn::{
    self, type">Module, type">Linear, LayerNorm, type">Dropout,
    MultiheadAttention, TransformerBlock, Embedding, type">Flatten,
};
use rumus::optim::{type">Trainer, type">AdamW};

"token-attribute">#[derive(type">Module)]
struct TinyTransformer {
    embedding: Embedding,
    block1: TransformerBlock,
    block2: TransformerBlock,
    ln_final: LayerNorm,
    fc_out: type">Linear,
}

impl TinyTransformer {
    fn new(vocab_size: type">usize, d_model: type">usize, n_heads: type">usize, ff_dim: type">usize) -> type">Self {
        type">Self {
            embedding: Embedding::new(vocab_size, d_model),
            block1: TransformerBlock::new(d_model, n_heads, ff_dim, 0.1),
            block2: TransformerBlock::new(d_model, n_heads, ff_dim, 0.1),
            ln_final: LayerNorm::new(d_model),
            fc_out: type">Linear::new(d_model, vocab_size),
        }
    }

    fn forward(&self, x: &type">Tensor) -> type">Tensor {
        let x = self.embedding.forward(x);        "token-comment">// [B, T] → [B, T, D]
        let x = self.block1.forward(&x);           "token-comment">// self-attn + FFN
        let x = self.block2.forward(&x);           "token-comment">// self-attn + FFN
        let x = self.ln_final.forward(&x);         "token-comment">// final layer norm
        "token-comment">// Take the last token's representation
        let last = x.select(1, -1);                "token-comment">// [B, D]
        self.fc_out.forward(&last)                 "token-comment">// [B, vocab]
    }
}

fn main() -> type">Result<(), Box<dyn std::error::Error>> {
    let mut model = TinyTransformer::new(10000, 128, 4, 512);
    model.to_gpu();  "token-comment">// move all parameters to GPU
    model.train();

    let optimizer = type">AdamW::new(model.parameters(), 3e-4);
    let mut trainer = type">Trainer::new(optimizer);

    for epoch in 0..20 {
        for (tokens, labels) in &train_loader {
            let tokens = tokens.to_gpu();
            let labels = labels.to_gpu();

            trainer.train_step(|| {
                let logits = model.forward(&tokens);
                nn::cross_entropy_loss(&logits, &labels)
            });
        }

        let avg_loss = trainer.epoch_avg_loss();
        println!("Epoch {}: loss = {:.4}", epoch, avg_loss);
    }

    "token-comment">// Save trained model
    model.eval();
    let state = model.state_dict("");
    nn::save_safetensors(&state, "transformer.safetensors")?;

    type">Ok(())
}