tract_core/ops/nn/
gelu_exact.rs

1use crate::internal::*;
2use crate::ops::binary::TypedBinOp;
3use crate::ops::math::{Add, Mul};
4
5const CHUNK: usize = 1024;
6
7fn inv_sqrt2() -> f32 {
8    (2.0f32).sqrt().recip()
9}
10
11crate::element_wise!(gelu_exact, GeluExact,
12    [f16] => |_, xs| {
13        let erf = (tract_linalg::ops().erf_f32)();
14        let c = f16::from_f32(inv_sqrt2());
15        let half = f16::from_f32(0.5);
16        let one = f16::from_f32(1.0);
17        let mut scratch = vec![0f32; xs.len().min(CHUNK)];
18        for chunk in xs.chunks_mut(CHUNK) {
19            let scaled = &mut scratch[..chunk.len()];
20            scaled.iter_mut().zip(chunk.iter()).for_each(|(s, x)| *s = (*x * c).to_f32());
21            erf.run(scaled)?;
22            chunk.iter_mut().zip(scaled.iter()).for_each(|(x, e)| {
23                *x = (*x * half) * (f16::from_f32(*e) + one);
24            });
25        }
26        Ok(())
27    },
28    [f32] => |_, xs| {
29        let erf = (tract_linalg::ops().erf_f32)();
30        let c = inv_sqrt2();
31        let mut scratch = vec![0f32; xs.len().min(CHUNK)];
32        for chunk in xs.chunks_mut(CHUNK) {
33            let scaled = &mut scratch[..chunk.len()];
34            scaled.iter_mut().zip(chunk.iter()).for_each(|(s, x)| *s = *x * c);
35            erf.run(scaled)?;
36            chunk.iter_mut().zip(scaled.iter()).for_each(|(x, e)| *x = (*x * 0.5) * (*e + 1.0));
37        }
38        Ok(())
39    };
40    cost: |dt| {tvec!((Cost::FMA(dt), 14), (Cost::Div(dt), 1))}
41);
42
43/// Search pattern => GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
44///
45/// Anchored on the `Erf`, which the ONNX `Gelu` (without `approximate="tanh"`)
46/// and `BiasGelu` expansions both emit as the middle of a five node chain.
47pub fn detect_gelu_exact(
48    model: &TypedModel,
49    node: &TypedNode,
50) -> TractResult<Option<TypedModelPatch>> {
51    let erf_node = node;
52    let dt = model.node_input_facts(erf_node.id)?[0].datum_type;
53    rule_if!(matches!(dt, DatumType::F32 | DatumType::F16));
54
55    // x / sqrt(2)
56    let scale = &model.nodes()[erf_node.inputs[0].node];
57    rule_if_some!(scale_op = scale.op_as::<TypedBinOp>());
58    rule_if!(scale_op.0.is::<Mul>());
59    rule_if!(model.matches_single_input_const(scale, inv_sqrt2()));
60    rule_if_some!(
61        x = scale
62            .inputs
63            .iter()
64            .find(|o| model.outlet_fact(**o).map(|f| f.konst.is_none()).unwrap_or(false))
65            .copied()
66    );
67
68    // 1 + erf(x / sqrt(2))
69    rule_if_some!(one_plus_erf = model.find_succ_bin_with_const::<Add>(erf_node, 1.0));
70
71    // (0.5 * x) * (1 + erf(x / sqrt(2)))
72    rule_if_some!(out = model.single_succ(one_plus_erf.id)?);
73    rule_if_some!(out_op = out.op_as::<TypedBinOp>());
74    rule_if!(out_op.0.is::<Mul>());
75    rule_if_some!(
76        half_x = out
77            .inputs
78            .iter()
79            .filter_map(|i| {
80                let n = &model.nodes()[i.node];
81                n.op_as::<TypedBinOp>()?.0.is::<Mul>().then_some(n)
82            })
83            .next()
84    );
85    rule_if!(model.matches_single_input_const(half_x, 0.5));
86    rule_if!(half_x.inputs.contains(&x));
87
88    let mut patch = TypedModelPatch::default();
89    let tap = patch.taps(model, &[x])?;
90    let wired =
91        patch.wire_node(format!("{}.gelu_exact", erf_node.name), gelu_exact(), &[tap[0]])?;
92    patch.shunt_outside(model, out.id.into(), wired[0])?;
93    Ok(Some(patch))
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::ops::element_wise::ElementWiseOp;
100    use crate::ops::math::{Erf, add, erf, mul};
101
102    fn chain(dt: DatumType, len: usize) -> TractResult<TypedModel> {
103        let mut m = TypedModel::default();
104        let x = m.add_source("x", dt.fact([len]))?;
105        let c = m.add_const("c", tensor1(&[inv_sqrt2()]).cast_to_dt(dt)?.into_owned())?;
106        let scaled = m.wire_node("scale", mul(), &[x, c])?[0];
107        let e = m.wire_node("erf", erf(), &[scaled])?[0];
108        let one = m.add_const("one", tensor1(&[1f32]).cast_to_dt(dt)?.into_owned())?;
109        let ope = m.wire_node("add_one", add(), &[e, one])?[0];
110        let half = m.add_const("half", tensor1(&[0.5f32]).cast_to_dt(dt)?.into_owned())?;
111        let hx = m.wire_node("half_x", mul(), &[x, half])?[0];
112        let out = m.wire_node("out", mul(), &[hx, ope])?;
113        m.select_output_outlets(&out)?;
114        Ok(m)
115    }
116
117    fn is_mini<T: crate::ops::element_wise::ElementWiseMiniOp>(n: &TypedNode) -> bool {
118        n.op_as::<ElementWiseOp>().map(|e| e.0.is::<T>()).unwrap_or(false)
119    }
120
121    fn input(dt: DatumType, len: usize) -> TractResult<TValue> {
122        let values: Vec<f32> = (0..len).map(|i| (i as f32 * 0.37).sin() * 4.0).collect();
123        Ok(tensor1(&values).cast_to_dt(dt)?.into_owned().into_tvalue())
124    }
125
126    #[test]
127    fn fuses_the_chain_and_keeps_the_values() -> TractResult<()> {
128        for dt in [DatumType::F32, DatumType::F16] {
129            let len = 2050;
130            let raw = chain(dt, len)?;
131            let reference = raw.clone().into_runnable()?.run(tvec!(input(dt, len)?))?;
132
133            let fused = raw.into_decluttered()?;
134            assert!(
135                fused.nodes().iter().any(is_mini::<GeluExact>),
136                "{dt:?}: no GeluExact after declutter"
137            );
138            assert!(!fused.nodes().iter().any(is_mini::<Erf>), "{dt:?}: Erf survived");
139
140            let got = fused.into_runnable()?.run(tvec!(input(dt, len)?))?;
141            assert_eq!(
142                got[0].as_bytes(),
143                reference[0].as_bytes(),
144                "{dt:?}: fused output differs from the chain"
145            );
146        }
147        Ok(())
148    }
149}
tract_core/ops/nn/gelu_exact.rs

tract_core/ops/nn/
gelu_exact.rs
