rlx_ir/ops/
normalization.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Normalization builders: LayerNorm, fused residual+LN (plan #53).
6
7use crate::{Graph, NodeId, Op, Shape};
8
9impl Graph {
10    /// LayerNorm2d on NCHW (normalize across channels at each spatial position).
11    pub fn layer_norm2d(&mut self, input: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId {
12        let shape = self.node(input).shape.clone();
13        self.push(
14            Op::LayerNorm2d { eps },
15            vec![input, gamma, beta],
16            shape,
17            None,
18        )
19    }
20
21    /// Group normalization on NCHW.
22    pub fn group_norm(
23        &mut self,
24        input: NodeId,
25        gamma: NodeId,
26        beta: NodeId,
27        num_groups: usize,
28        eps: f32,
29    ) -> NodeId {
30        let shape = self.node(input).shape.clone();
31        self.push(
32            Op::GroupNorm { num_groups, eps },
33            vec![input, gamma, beta],
34            shape,
35            None,
36        )
37    }
38
39    /// BatchNorm inference (frozen running mean/variance).
40    pub fn batch_norm_inference(
41        &mut self,
42        input: NodeId,
43        gamma: NodeId,
44        beta: NodeId,
45        running_mean: NodeId,
46        running_var: NodeId,
47        eps: f32,
48    ) -> NodeId {
49        let shape = self.node(input).shape.clone();
50        self.push(
51            Op::BatchNormInference { eps },
52            vec![input, gamma, beta, running_mean, running_var],
53            shape,
54            None,
55        )
56    }
57
58    /// Layer normalization.
59    pub fn layer_norm(
60        &mut self,
61        input: NodeId,
62        gamma: NodeId,
63        beta: NodeId,
64        axis: i32,
65        eps: f32,
66        shape: Shape,
67    ) -> NodeId {
68        self.push(
69            Op::LayerNorm { axis, eps },
70            vec![input, gamma, beta],
71            shape,
72            None,
73        )
74    }
75
76    /// Fused residual + bias + layer norm (created by optimization passes).
77    pub fn fused_residual_ln(
78        &mut self,
79        x: NodeId,
80        residual: NodeId,
81        bias: Option<NodeId>,
82        gamma: NodeId,
83        beta: NodeId,
84        eps: f32,
85        shape: Shape,
86    ) -> NodeId {
87        let has_bias = bias.is_some();
88        let mut inputs = vec![x, residual];
89        if let Some(b) = bias {
90            inputs.push(b);
91        }
92        inputs.push(gamma);
93        inputs.push(beta);
94        self.push(Op::FusedResidualLN { has_bias, eps }, inputs, shape, None)
95    }
96
97    /// Fused residual + bias + RMS norm (created by optimization passes).
98    pub fn fused_residual_rms_norm(
99        &mut self,
100        x: NodeId,
101        residual: NodeId,
102        bias: Option<NodeId>,
103        gamma: NodeId,
104        beta: NodeId,
105        eps: f32,
106        shape: Shape,
107    ) -> NodeId {
108        let has_bias = bias.is_some();
109        let mut inputs = vec![x, residual];
110        if let Some(b) = bias {
111            inputs.push(b);
112        }
113        inputs.push(gamma);
114        inputs.push(beta);
115        self.push(
116            Op::FusedResidualRmsNorm { has_bias, eps },
117            inputs,
118            shape,
119            None,
120        )
121    }
122}
rlx_ir/ops/normalization.rs

rlx_ir/ops/
normalization.rs
