vyre_libs/nn/norm/
last_dim_l2_norm.rs

1//! Last-dimension L2 normalization with float32 accumulation.
2
3use thiserror::Error;
4use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program, UnOp};
5
6use crate::region::wrap_anonymous;
7
8const OP_ID: &str = "vyre-libs::nn::last_dim_l2_norm";
9
10/// Invalid last-dimension L2 normalization construction.
11#[derive(Debug, Clone, PartialEq, Eq, Error)]
12pub enum LastDimL2NormError {
13    /// A tensor dimension is zero.
14    #[error("last-dimension L2 normalization requires nonzero rows and width; got rows={rows}, width={width}")]
15    EmptyShape {
16        /// Row count.
17        rows: u32,
18        /// Last-dimension width.
19        width: u32,
20    },
21    /// Flattened element count exceeds u32 indexing.
22    #[error("last-dimension L2 normalization rows*width overflows u32; split the tensor")]
23    ElementCountOverflow,
24    /// Source dtype lacks the required floating conversion contract.
25    #[error("last-dimension L2 normalization supports F16, BF16, or F32 tensors; got {dtype:?}")]
26    UnsupportedDtype {
27        /// Rejected dtype.
28        dtype: DataType,
29    },
30}
31
32/// Build `output = input * rsqrt(sum(input², last_dim) + eps)`.
33///
34/// Every row accumulates in F32. The normalized result converts once to the
35/// source dtype at the output boundary.
36///
37/// # Errors
38///
39/// Returns [`LastDimL2NormError`] for empty or overflowing shapes and for
40/// source dtypes without F16, BF16, or F32 conversion semantics.
41pub fn last_dim_l2_norm(
42    input: &str,
43    output: &str,
44    rows: u32,
45    width: u32,
46    eps: f32,
47    dtype: DataType,
48) -> Result<Program, LastDimL2NormError> {
49    if rows == 0 || width == 0 {
50        return Err(LastDimL2NormError::EmptyShape { rows, width });
51    }
52    if !matches!(dtype, DataType::F16 | DataType::BF16 | DataType::F32) {
53        return Err(LastDimL2NormError::UnsupportedDtype { dtype });
54    }
55    let total = rows
56        .checked_mul(width)
57        .ok_or(LastDimL2NormError::ElementCountOverflow)?;
58    let index = Expr::var("index");
59    let row_start = Expr::mul(Expr::div(index.clone(), Expr::u32(width)), Expr::u32(width));
60    let normalized = Expr::mul(
61        Expr::cast(DataType::F32, Expr::load(input, index.clone())),
62        Expr::UnOp {
63            op: UnOp::InverseSqrt,
64            operand: Box::new(Expr::add(Expr::var("sum_squares"), Expr::f32(eps))),
65        },
66    );
67    let body = vec![
68        Node::let_bind("index", Expr::InvocationId { axis: 0 }),
69        Node::if_then(
70            Expr::lt(index.clone(), Expr::u32(total)),
71            vec![
72                Node::let_bind("row_start", row_start),
73                Node::let_bind("sum_squares", Expr::f32(0.0)),
74                Node::loop_for(
75                    "offset",
76                    Expr::u32(0),
77                    Expr::u32(width),
78                    vec![
79                        Node::let_bind(
80                            "l2_value",
81                            Expr::cast(
82                                DataType::F32,
83                                Expr::load(
84                                    input,
85                                    Expr::add(Expr::var("row_start"), Expr::var("offset")),
86                                ),
87                            ),
88                        ),
89                        Node::assign(
90                            "sum_squares",
91                            Expr::add(
92                                Expr::var("sum_squares"),
93                                Expr::mul(Expr::var("l2_value"), Expr::var("l2_value")),
94                            ),
95                        ),
96                    ],
97                ),
98                Node::Store {
99                    buffer: output.into(),
100                    index,
101                    value: Expr::cast(dtype.clone(), normalized),
102                },
103            ],
104        ),
105    ];
106    Ok(Program::wrapped(
107        vec![
108            BufferDecl::storage(input, 0, BufferAccess::ReadOnly, dtype.clone()).with_count(total),
109            BufferDecl::output(output, 1, dtype).with_count(total),
110        ],
111        [64, 1, 1],
112        vec![wrap_anonymous(OP_ID, body)],
113    ))
114}
vyre_libs/nn/norm/last_dim_l2_norm.rs

vyre_libs/nn/norm/
last_dim_l2_norm.rs
