tract_core/ops/nn/
gelu_exact.rs

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

tract_core/ops/nn/
gelu_exact.rs
