tract_core/ops/nn/
gelu_approximate.rs

1use crate::internal::*;
2use crate::ops::binary::TypedBinOp;
3use crate::ops::element_wise::ElementWiseOp;
4use crate::ops::math::{Add, Mul, Pow, Tanh};
5
6use tract_data::half::f16;
7use tract_linalg::routines::Func;
8
9fn gelu_approx_f32(x: f32, pow: i32) -> f32 {
10    let sqrt_2_over_pi = (2.0 / std::f32::consts::PI).sqrt();
11    0.5 * x * (1.0 + f32::tanh(sqrt_2_over_pi * (x + 0.044715 * x.powi(pow))))
12}
13
14element_wise!(gelu_approximate, GeluApproximate { fast_impl: bool },
15    [f16] => |op, xs| {
16        let pow = if op.fast_impl { 2 } else { 3 };
17        xs.iter_mut().for_each(|x| {
18            *x = f16::from_f32(gelu_approx_f32(x.to_f32(), pow));
19        });
20        Ok(())
21    },
22    [f32] => |op, xs| {
23        if op.fast_impl {
24            // pow=2 fast path: no linalg kernel yet, scalar fallback.
25            xs.iter_mut().for_each(|x| {
26                *x = gelu_approx_f32(*x, 2);
27            });
28            Ok(())
29        } else {
30            // pow=3 canonical path: linalg NEON kernel composes with tanh.
31            Func::Gelu.ew_f32()?.run(xs)
32        }
33    };
34    cost: |dt| {tvec!((Cost::FMA(dt), 15))}
35);
36
37/// Search pattern => NEW_GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^N))); N ∈ {2, 3}
38pub fn detect_gelu_approx(
39    _op: &Pow,
40    model: &TypedModel,
41    node: &TypedNode,
42) -> TractResult<Option<TypedModelPatch>> {
43    let pow_node = node;
44
45    let in_fact = model.node_input_facts(pow_node.id)?[0];
46    let dt = in_fact.datum_type;
47
48    // Only F16 and F32 is supported.
49    rule_if!(matches!(dt, DatumType::F32 | DatumType::F16));
50
51    rule_if!(
52        model.matches_single_input_const(pow_node, 3.0)
53            || model.matches_single_input_const(pow_node, 2.0)
54    );
55    let fast_impl = model.matches_single_input_const(pow_node, 2.0);
56
57    // 0.044715 * x^N
58    rule_if_some!(mul_coef_a = model.find_succ_bin_with_const::<Mul>(pow_node, 0.044715));
59
60    // x + 0.044715 * x^N
61    rule_if_some!(
62        x_plus_mul_coef_a = model.find_succ_bin_with_outlet::<Add>(mul_coef_a, &pow_node.inputs[0])
63    );
64
65    // sqrt(2/pi) * (x + 0.044715 * x^N)
66    let sqrt_2_over_pi = (2.0 / std::f32::consts::PI).sqrt();
67    rule_if_some!(
68        mul_sqrt_2_over_pi =
69            model.find_succ_bin_with_const::<Mul>(x_plus_mul_coef_a, sqrt_2_over_pi)
70    );
71
72    // tanh(sqrt(2/pi) * (x + 0.044715 * x^N))
73    rule_if_some!(tanh_succ = model.single_succ(mul_sqrt_2_over_pi.id)?);
74    rule_if_some!(tanh_succ_op = tanh_succ.op_as::<ElementWiseOp>());
75    rule_if!(tanh_succ_op.0.is::<Tanh>());
76
77    // 1.0 + tanh(sqrt(2/pi) * (x + 0.044715 * x^N)) N ∈ {2, 3}
78    rule_if_some!(tanh_plus_1 = model.find_succ_bin_with_const::<Add>(tanh_succ, 1.0));
79
80    // Identify Mul
81    rule_if_some!(mul_succ = model.single_succ(tanh_plus_1.id)?);
82    rule_if_some!(mul_succ_op = mul_succ.op_as::<TypedBinOp>());
83    rule_if!(mul_succ_op.0.is::<Mul>());
84
85    // Search first
86    // tmp = x * (1.0 + tanh(sqrt(2/pi) * (x + 0.044715 * x^N)))
87    // out = 0.5 * tmp
88    let last_node_id = if mul_succ.inputs.contains(&pow_node.inputs[0]) {
89        // 0.5 * x * (1.0 + tanh(sqrt(2/pi) * (x + 0.044715 * x^N)))
90        rule_if_some!(last_mul_with_0_5 = model.find_succ_bin_with_const::<Mul>(mul_succ, 0.5));
91        last_mul_with_0_5.id
92    } else {
93        // tmp = 0.5 * x
94        // out = tmp * (1.0 + tanh(sqrt(2/pi) * (x + 0.044715 * x^N))) N ∈ {2, 3}
95        rule_if_some!(
96            x_mul_0_5 = mul_succ
97                .inputs
98                .iter()
99                .filter_map(|i| {
100                    let n = &model.nodes()[i.node];
101                    let op = n.op_as::<TypedBinOp>()?;
102                    op.0.is::<Mul>().then_some(n)
103                })
104                .next()
105        );
106        rule_if!(model.matches_single_input_const(x_mul_0_5, 0.5));
107        rule_if!(x_mul_0_5.inputs.contains(&pow_node.inputs[0]));
108        mul_succ.id
109    };
110
111    let mut patch = TypedModelPatch::default();
112    let gelu_approx_input = patch.taps(model, &pow_node.inputs)?;
113    let out = patch.wire_node(
114        format!("{}.gelu_approx", pow_node.name),
115        gelu_approximate(fast_impl),
116        &[gelu_approx_input[0]],
117    )?;
118    patch.shunt_outside(model, last_node_id.into(), out[0])?;
119    Ok(Some(patch))
120}
tract_core/ops/nn/gelu_approximate.rs

tract_core/ops/nn/
gelu_approximate.rs
