foxhunt/crates/ml-supervised/src/diffusion/denoiser.rs

//! Denoiser network for the diffusion model.
//!
//! Uses a fully-connected architecture with sinusoidal time embedding.
//! All computation uses cuBLAS-backed `GpuLinear` layers.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{gpu_add, gpu_silu, GpuLinear, GpuTensor};

/// Sinusoidal time embedding for diffusion timestep conditioning.
///
/// Maps scalar timestep t to a fixed-dimension vector using
/// sin/cos positional encoding (same idea as Transformer PE).
pub struct TimeEmbedding {
    proj: GpuLinear,
    embed_dim: usize,
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for TimeEmbedding {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TimeEmbedding").finish_non_exhaustive()
    }
}

impl TimeEmbedding {
    pub fn new(
        embed_dim: usize,
        hidden_dim: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let proj = GpuLinear::new(embed_dim, hidden_dim, stream)?;
        Ok(Self {
            proj,
            embed_dim,
            stream: Arc::clone(stream),
        })
    }

    /// Compute sinusoidal embedding for timestep indices.
    /// Input: Vec<u32> timestep indices -> Output: (batch, `hidden_dim`)
    pub fn forward(&self, t: &[u32]) -> Result<GpuTensor, MLError> {
        let half_dim = self.embed_dim / 2;
        let batch = t.len();

        let mut freq_vals = Vec::with_capacity(half_dim);
        for i in 0..half_dim {
            let freq = (-(10000.0_f64.ln()) * i as f64 / half_dim.max(1) as f64).exp();
            freq_vals.push(freq as f32);
        }

        // Compute sinusoidal embedding on CPU
        let mut emb = vec![0.0_f32; batch * self.embed_dim];
        for (b, &t_val) in t.iter().enumerate() {
            for (i, &freq) in freq_vals.iter().enumerate() {
                let angle = t_val as f32 * freq;
                // sin half
                if let Some(slot) = emb.get_mut(b * self.embed_dim + i) {
                    *slot = angle.sin();
                }
                // cos half
                if let Some(slot) = emb.get_mut(b * self.embed_dim + half_dim + i) {
                    *slot = angle.cos();
                }
            }
        }

        let emb_tensor = GpuTensor::from_vec(emb, &[batch, self.embed_dim], &self.stream)?;

        // Project to hidden_dim
        self.proj.forward(&emb_tensor)
    }
}

/// A single denoiser block: linear -> `SiLU` -> linear + time conditioning + residual.
struct DenoiserBlock {
    fc1: GpuLinear,
    fc2: GpuLinear,
    time_proj: GpuLinear,
    has_residual: bool,
}

impl std::fmt::Debug for DenoiserBlock {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DenoiserBlock").finish_non_exhaustive()
    }
}

impl DenoiserBlock {
    fn new(
        input_dim: usize,
        hidden_dim: usize,
        time_dim: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let fc1 = GpuLinear::new(input_dim, hidden_dim, stream)?;
        let fc2 = GpuLinear::new(hidden_dim, hidden_dim, stream)?;
        let time_proj = GpuLinear::new(time_dim, hidden_dim, stream)?;
        Ok(Self {
            fc1,
            fc2,
            time_proj,
            has_residual: input_dim == hidden_dim,
        })
    }

    fn forward(&self, x: &GpuTensor, t_emb: &GpuTensor) -> Result<GpuTensor, MLError> {
        // fc1 -> SiLU
        let h = self.fc1.forward(x)?;
        let h = gpu_silu(&h)?;

        // Add time embedding
        let t_proj = self.time_proj.forward(t_emb)?;
        let h = gpu_add(&h, &t_proj)?;

        // fc2 -> SiLU
        let h = self.fc2.forward(&h)?;
        let h = gpu_silu(&h)?;

        // Residual connection
        if self.has_residual {
            gpu_add(&h, x)
        } else {
            Ok(h)
        }
    }
}

/// Fully-connected denoiser network for diffusion.
///
/// Architecture: input projection -> N denoiser blocks -> output projection.
/// Each block is time-conditioned via additive time embedding.
pub struct Denoiser {
    input_proj: GpuLinear,
    blocks: Vec<DenoiserBlock>,
    output_proj: GpuLinear,
    time_embed: TimeEmbedding,
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for Denoiser {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Denoiser").finish_non_exhaustive()
    }
}

impl Denoiser {
    pub fn new(
        data_dim: usize,
        hidden_dim: usize,
        num_layers: usize,
        time_embed_dim: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        if data_dim == 0 || hidden_dim == 0 {
            return Err(MLError::ConfigError(format!(
                "Diffusion denoiser requires data_dim > 0 and hidden_dim > 0 (got {}x{})",
                data_dim, hidden_dim
            )));
        }

        let time_embed = TimeEmbedding::new(time_embed_dim, hidden_dim, stream)?;
        let input_proj = GpuLinear::new(data_dim, hidden_dim, stream)?;

        let mut blocks = Vec::with_capacity(num_layers);
        for _i in 0..num_layers {
            let block = DenoiserBlock::new(hidden_dim, hidden_dim, hidden_dim, stream)?;
            blocks.push(block);
        }

        let output_proj = GpuLinear::new(hidden_dim, data_dim, stream)?;

        Ok(Self {
            input_proj,
            blocks,
            output_proj,
            time_embed,
            stream: Arc::clone(stream),
        })
    }

    /// Predict noise epsilon given noisy input `x_t` and timestep t.
    ///
    /// Input x: (batch, `data_dim`), t: &[u32] -> Output: (batch, `data_dim`)
    pub fn forward(&self, x: &GpuTensor, t: &[u32]) -> Result<GpuTensor, MLError> {
        // Time embedding
        let t_emb = self.time_embed.forward(t)?;

        // Input projection -> SiLU
        let mut h = self.input_proj.forward(x)?;
        h = gpu_silu(&h)?;

        // Denoiser blocks
        for block in &self.blocks {
            h = block.forward(&h, &t_emb)?;
        }

        // Output projection -> predicted noise
        self.output_proj.forward(&h)
    }

    /// Get a reference to the stream.
    pub const fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Collect all `GpuLinear` layers for weight access.
    pub fn all_linears(&self) -> Vec<&GpuLinear> {
        let mut linears = vec![
            &self.input_proj,
            &self.time_embed.proj,
        ];
        for block in &self.blocks {
            linears.push(&block.fc1);
            linears.push(&block.fc2);
            linears.push(&block.time_proj);
        }
        linears.push(&self.output_proj);
        linears
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA required");
        ctx.new_stream().expect("Failed to create stream")
    }

    #[test]
    fn test_time_embedding_shape() {
        let stream = test_stream();
        let te = TimeEmbedding::new(32, 64, &stream).unwrap();
        let t = vec![0_u32, 100, 500, 999];
        let emb = te.forward(&t).unwrap();
        assert_eq!(emb.shape, vec![4, 64]);
    }

    #[test]
    fn test_time_embedding_different_timesteps_differ() {
        let stream = test_stream();
        let te = TimeEmbedding::new(32, 64, &stream).unwrap();
        let e1 = te.forward(&[0_u32]).unwrap().to_vec().unwrap(); // test-only readback
        let e2 = te.forward(&[500_u32]).unwrap().to_vec().unwrap(); // test-only readback
        let diff: f32 = e1.iter().zip(e2.iter()).map(|(a, b)| (a - b).abs()).sum();
        assert!(diff > 0.0, "Different timesteps should produce different embeddings");
    }

    #[test]
    fn test_denoiser_output_shape() {
        let stream = test_stream();
        let denoiser = Denoiser::new(64, 128, 3, 32, &stream).unwrap();
        let x = GpuTensor::randn(&[4, 64], 1.0, &stream).unwrap();
        let t = vec![100_u32, 200, 300, 400];
        let out = denoiser.forward(&x, &t).unwrap();
        assert_eq!(out.shape, vec![4, 64]);
    }

    #[test]
    fn test_denoiser_single_layer() {
        let stream = test_stream();
        let denoiser = Denoiser::new(32, 64, 1, 16, &stream).unwrap();
        let x = GpuTensor::randn(&[2, 32], 1.0, &stream).unwrap();
        let t = vec![0_u32, 999];
        let out = denoiser.forward(&x, &t).unwrap();
        assert_eq!(out.shape, vec![2, 32]);
    }

    #[test]
    fn test_denoiser_finite_output() {
        let stream = test_stream();
        let denoiser = Denoiser::new(16, 32, 2, 8, &stream).unwrap();
        let x = GpuTensor::randn(&[2, 16], 0.5, &stream).unwrap();
        let t = vec![50_u32, 100];
        let out = denoiser.forward(&x, &t).unwrap();
        let v = out.to_vec().unwrap(); // test-only readback
        for val in &v {
            assert!(val.is_finite(), "Non-finite denoiser output: {}", val);
        }
    }
}