himada-integrations 0.1.0

use candle_core::{CustomOp1, CustomOp2, Result, Tensor};

use himada_dispatch::*;

thread_local! {
    static DOT: Dispatch<DotKernel> = Dispatch::new("candle_dot", vec![]);
    static SOFTMAX: Dispatch<SoftmaxKernel> = Dispatch::new("candle_softmax", vec![]);
    static NEGATE: Dispatch<SoftmaxKernel> = Dispatch::new("candle_negate", vec![]);
    static CLAMP: Dispatch<ClampKernel> = Dispatch::new("candle_clamp", vec![]);
}

pub struct HimadaDot;

impl CustomOp2 for HimadaDot {
    fn name(&self) -> &'static str {
        "himada-dot"
    }

    fn cpu_fwd(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
        let a_s: Vec<f64> = a.to_vec1()?;
        let b_s: Vec<f64> = b.to_vec1()?;
        let result = DOT.with(|d| d.compute(&a_s, &b_s));
        Tensor::new(result, a.device())
    }
}

pub struct HimadaSoftmax;

impl CustomOp1 for HimadaSoftmax {
    fn name(&self) -> &'static str {
        "himada-softmax"
    }

    fn cpu_fwd(&self, a: &Tensor) -> Result<Tensor> {
        let a_s: Vec<f64> = a.to_vec1()?;
        let mut out = vec![0.0; a_s.len()];
        SOFTMAX.with(|d| d.compute(&a_s, &mut out));
        Tensor::from_vec(out, a.shape(), a.device())
    }
}

pub struct HimadaClamp {
    lo: f64,
    hi: f64,
}

impl HimadaClamp {
    pub fn new(lo: f64, hi: f64) -> Self {
        Self { lo, hi }
    }
}

impl CustomOp1 for HimadaClamp {
    fn name(&self) -> &'static str {
        "himada-clamp"
    }

    fn cpu_fwd(&self, a: &Tensor) -> Result<Tensor> {
        let a_s: Vec<f64> = a.to_vec1()?;
        let mut out = vec![0.0; a_s.len()];
        CLAMP.with(|d| d.compute(&a_s, self.lo, self.hi, &mut out));
        Tensor::from_vec(out, a.shape(), a.device())
    }
}