onnx-runtime-ep-cpu 0.1.0-dev.0

//! `LayerNormalization`: normalize f32 `X` over the axes from `axis` onward,
//! then scale and shift (`docs/ORT2.md` §4.4).
//!
//! `Y = (X - mean) / sqrt(var + epsilon) * Scale + B`, where `mean`/`var` are
//! the population statistics over the normalized axes. The optional `Mean` and
//! `InvStdDev` outputs are filled when the caller provides those output slots.

use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, TensorView};
use onnx_runtime_ir::Node;

use super::{check_arity, to_dense_f32, write_dense_f32};

/// f32 LayerNormalization kernel carrying `axis` and `epsilon`.
pub struct LayerNormKernel {
    axis: i64,
    epsilon: f32,
}

/// Factory reading `axis` (default -1) and `epsilon` (default 1e-5).
pub struct LayerNormFactory;

impl KernelFactory for LayerNormFactory {
    fn create(&self, node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        let axis = node.attr("axis").and_then(|a| a.as_int()).unwrap_or(-1);
        let epsilon = node
            .attr("epsilon")
            .and_then(|a| a.as_float())
            .unwrap_or(1e-5);
        Ok(Box::new(LayerNormKernel { axis, epsilon }))
    }
}

impl Kernel for LayerNormKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        check_arity("LayerNormalization", inputs, outputs, 2, 3, 1)?;
        let x = to_dense_f32(&inputs[0])?;
        let scale = to_dense_f32(&inputs[1])?;
        let bias = if inputs.len() == 3 {
            Some(to_dense_f32(&inputs[2])?)
        } else {
            None
        };

        let x_shape = inputs[0].shape;
        let rank = x_shape.len();
        let axis = if self.axis < 0 {
            self.axis + rank as i64
        } else {
            self.axis
        };
        if axis < 0 || axis as usize > rank {
            return Err(EpError::KernelFailed(format!(
                "LayerNormalization: axis {} out of range for rank {rank}",
                self.axis
            )));
        }
        let axis = axis as usize;

        let norm_size: usize = x_shape[axis..].iter().product();
        let num_groups: usize = x_shape[..axis].iter().product();
        if norm_size == 0 {
            return Err(EpError::KernelFailed(
                "LayerNormalization: empty normalization axis".into(),
            ));
        }
        if scale.len() != norm_size {
            return Err(EpError::KernelFailed(format!(
                "LayerNormalization: scale has {} elements, expected {norm_size}",
                scale.len()
            )));
        }
        if let Some(b) = &bias
            && b.len() != norm_size
        {
            return Err(EpError::KernelFailed(format!(
                "LayerNormalization: bias has {} elements, expected {norm_size}",
                b.len()
            )));
        }
        let mut y = vec![0.0f32; x.len()];
        let mut means = vec![0.0f32; num_groups];
        let mut inv_stds = vec![0.0f32; num_groups];

        for g in 0..num_groups {
            let base = g * norm_size;
            let slice = &x[base..base + norm_size];
            let mean = slice.iter().sum::<f32>() / norm_size as f32;
            let var =
                slice.iter().map(|&v| (v - mean) * (v - mean)).sum::<f32>() / norm_size as f32;
            let inv_std = 1.0 / (var + self.epsilon).sqrt();
            means[g] = mean;
            inv_stds[g] = inv_std;
            for e in 0..norm_size {
                let norm = (slice[e] - mean) * inv_std;
                let mut out = norm * scale[e];
                if let Some(b) = &bias {
                    out += b[e];
                }
                y[base + e] = out;
            }
        }

        write_dense_f32(&mut outputs[0], &y)?;
        // Optional Mean / InvStdDev outputs (shape = X.shape[:axis] with the
        // normalized axes as 1s; element count == num_groups).
        if outputs.len() > 1 {
            write_dense_f32(&mut outputs[1], &means)?;
        }
        if outputs.len() > 2 {
            write_dense_f32(&mut outputs[2], &inv_stds)?;
        }
        Ok(())
    }

    fn supports_strided_input(&self, _input_idx: usize) -> bool {
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::kernels::testutil::Owned;

    #[test]
    fn layernorm_last_axis_matches_reference() {
        // X [2,3], normalize last axis. Scale=1, Bias=0, eps=1e-5.
        let x = Owned::f32(&[2, 3], &[1., 2., 3., 2., 4., 6.]);
        let scale = Owned::f32(&[3], &[1., 1., 1.]);
        let mut out = Owned::zeros_f32(&[2, 3]);
        let k = LayerNormKernel {
            axis: -1,
            epsilon: 1e-5,
        };
        k.execute(&[x.view(), scale.view()], &mut [out.view_mut()])
            .unwrap();

        // Hand-computed: row0 mean=2, var=2/3, inv=1/sqrt(0.66667+1e-5)
        let expect = |row: &[f32]| -> Vec<f32> {
            let mean = row.iter().sum::<f32>() / 3.0;
            let var = row.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / 3.0;
            let inv = 1.0 / (var + 1e-5f32).sqrt();
            row.iter().map(|v| (v - mean) * inv).collect()
        };
        let mut want = expect(&[1., 2., 3.]);
        want.extend(expect(&[2., 4., 6.]));
        let got = out.to_f32();
        for (g, w) in got.iter().zip(&want) {
            assert!((g - w).abs() < 1e-5, "got {g}, want {w}");
        }
    }

    #[test]
    fn layernorm_applies_scale_and_bias() {
        let x = Owned::f32(&[1, 4], &[1., 2., 3., 4.]);
        let scale = Owned::f32(&[4], &[2., 2., 2., 2.]);
        let bias = Owned::f32(&[4], &[1., 1., 1., 1.]);
        let mut out = Owned::zeros_f32(&[1, 4]);
        let k = LayerNormKernel {
            axis: -1,
            epsilon: 1e-5,
        };
        k.execute(
            &[x.view(), scale.view(), bias.view()],
            &mut [out.view_mut()],
        )
        .unwrap();

        let row = [1., 2., 3., 4.];
        let mean = 2.5;
        let var = row.iter().map(|v: &f32| (v - mean).powi(2)).sum::<f32>() / 4.0;
        let inv = 1.0 / (var + 1e-5f32).sqrt();
        let want: Vec<f32> = row.iter().map(|v| (v - mean) * inv * 2.0 + 1.0).collect();
        for (g, w) in out.to_f32().iter().zip(&want) {
            assert!((g - w).abs() < 1e-5, "got {g}, want {w}");
        }
    }

    #[test]
    fn layernorm_writes_optional_mean_and_invstd() {
        let x = Owned::f32(&[2, 3], &[1., 2., 3., 2., 4., 6.]);
        let scale = Owned::f32(&[3], &[1., 1., 1.]);
        let mut y = Owned::zeros_f32(&[2, 3]);
        let mut mean = Owned::zeros_f32(&[2, 1]);
        let mut invstd = Owned::zeros_f32(&[2, 1]);
        let k = LayerNormKernel {
            axis: -1,
            epsilon: 1e-5,
        };
        k.execute(
            &[x.view(), scale.view()],
            &mut [y.view_mut(), mean.view_mut(), invstd.view_mut()],
        )
        .unwrap();
        assert!((mean.to_f32()[0] - 2.0).abs() < 1e-6);
        assert!((mean.to_f32()[1] - 4.0).abs() < 1e-6);
    }
}