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

tract_core/ops/nn/
gelu_approximate.rs
