teeny_kernels/nn/activation/
gelu.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#![allow(non_snake_case)]
18
19use teeny_core::dtype::Float;
20use teeny_macros::kernel;
21use teeny_triton::triton::{
22    types::{AddOffsets, Comparison},
23    *,
24};
25
26// ── GELU ─────────────────────────────────────────────────────────────────────
27
28/// Forward: y = x / (1 + exp(-2 * c * (x + a*x³)))
29///   where c = sqrt(2/pi), a = 0.044715 — the tanh GELU approximation.
30// ANCHOR: gelu_forward
31#[kernel(backward = GeluBackward)]
32pub fn gelu_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
33    x_ptr: T::Pointer<D>,
34    y_ptr: T::Pointer<D>,
35    n_elements: i32,
36) where
37    T::I32Tensor: types::Tensor<i32, 1>,
38    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
39    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
40{
41    let pid = T::program_id(Axis::X);
42    let block_start = pid * BLOCK_SIZE;
43    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
44    let in_bounds = offsets.lt(n_elements);
45
46    let x = T::load(
47        x_ptr.add_offsets(offsets),
48        Some(in_bounds),
49        None,
50        &[],
51        None,
52        None,
53        None,
54        false,
55    );
56
57    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
58    let neg2c = T::full(&[BLOCK_SIZE], D::from_f64(-2.0 * 0.7978845608028654));
59    let coeff = T::full(&[BLOCK_SIZE], D::from_f64(0.044715));
60
61    // tanh-GELU: y = x * 0.5 * (1 + tanh(c*(x + a*x³)))
62    //              = x / (1 + exp(-2c*(x + a*x³)))
63    let inner = x + coeff * x * x * x;
64    let y = x / (one + T::exp(neg2c * inner));
65    T::store(
66        y_ptr.add_offsets(offsets),
67        y,
68        Some(in_bounds),
69        &[],
70        None,
71        None,
72    );
73}
74
75// ANCHOR_END: gelu_forward
76
77/// Backward of the tanh-GELU approximation.
78///   Let inner = x + a*x³, s = sigmoid(2c*inner), t = tanh(c*inner) = 2s-1
79///   d/dx = 0.5*(1 + t) + x * 0.5 * sech²(c*inner) * c*(1+3a*x²)
80#[kernel]
81pub fn gelu_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
82    dy_ptr: T::Pointer<D>,
83    x_ptr: T::Pointer<D>,
84    dx_ptr: T::Pointer<D>,
85    n_elements: i32,
86) where
87    T::I32Tensor: types::Tensor<i32, 1>,
88    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
89    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
90{
91    let pid = T::program_id(Axis::X);
92    let block_start = pid * BLOCK_SIZE;
93    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
94    let in_bounds = offsets.lt(n_elements);
95
96    let dy = T::load(
97        dy_ptr.add_offsets(offsets),
98        Some(in_bounds),
99        None,
100        &[],
101        None,
102        None,
103        None,
104        false,
105    );
106    let x = T::load(
107        x_ptr.add_offsets(offsets),
108        Some(in_bounds),
109        None,
110        &[],
111        None,
112        None,
113        None,
114        false,
115    );
116
117    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
118    let half = T::full(&[BLOCK_SIZE], D::from_f64(0.5));
119    let two = T::full(&[BLOCK_SIZE], D::from_f64(2.0));
120    let three = T::full(&[BLOCK_SIZE], D::from_f64(3.0));
121    let c = T::full(&[BLOCK_SIZE], D::from_f64(0.7978845608028654));
122    let neg2c = T::full(&[BLOCK_SIZE], D::from_f64(-2.0 * 0.7978845608028654));
123    let coeff = T::full(&[BLOCK_SIZE], D::from_f64(0.044715));
124
125    let inner = x + coeff * x * x * x;
126    let s = one / (one + T::exp(neg2c * inner)); // sigmoid(2c*inner)
127    let t = two * s - one; // tanh(c*inner)
128    let sech2 = one - t * t; // 1 - tanh²
129    let dinner = c * (one + three * coeff * x * x);
130    let dx = dy * (half * (one + t) + x * half * sech2 * dinner);
131    T::store(
132        dx_ptr.add_offsets(offsets),
133        dx,
134        Some(in_bounds),
135        &[],
136        None,
137        None,
138    );
139}
140
141// ── Mish ─────────────────────────────────────────────────────────────────────
142
143/// Forward: y = x * tanh(softplus(x)) = x * tanh(log(1 + exp(x)))
144#[kernel(backward = MishBackward)]
145pub fn mish_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
146    x_ptr: T::Pointer<D>,
147    y_ptr: T::Pointer<D>,
148    n_elements: i32,
149) where
150    T::I32Tensor: types::Tensor<i32, 1>,
151    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
152    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
153{
154    let pid = T::program_id(Axis::X);
155    let block_start = pid * BLOCK_SIZE;
156    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
157    let in_bounds = offsets.lt(n_elements);
158
159    let x = T::load(
160        x_ptr.add_offsets(offsets),
161        Some(in_bounds),
162        None,
163        &[],
164        None,
165        None,
166        None,
167        false,
168    );
169    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
170    let two = T::full(&[BLOCK_SIZE], D::from_f64(2.0));
171    let neg2 = T::full(&[BLOCK_SIZE], D::from_f64(-2.0));
172    let sp = T::log(one + T::exp(x)); // softplus(x)
173    // tanh(sp) = 2*sigmoid(2*sp) - 1 = 2/(1+exp(-2*sp)) - 1
174    let s2 = one / (one + T::exp(neg2 * sp));
175    let t = two * s2 - one;
176    let y = x * t;
177    T::store(
178        y_ptr.add_offsets(offsets),
179        y,
180        Some(in_bounds),
181        &[],
182        None,
183        None,
184    );
185}
186
187/// Backward: dx = dy * (tanh(sp) + x * (1 - tanh²(sp)) * sigmoid(x))
188///   where sp = softplus(x). Recomputes all intermediates from x.
189#[kernel]
190pub fn mish_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
191    dy_ptr: T::Pointer<D>,
192    x_ptr: T::Pointer<D>,
193    dx_ptr: T::Pointer<D>,
194    n_elements: i32,
195) where
196    T::I32Tensor: types::Tensor<i32, 1>,
197    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
198    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
199{
200    let pid = T::program_id(Axis::X);
201    let block_start = pid * BLOCK_SIZE;
202    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
203    let in_bounds = offsets.lt(n_elements);
204
205    let dy = T::load(
206        dy_ptr.add_offsets(offsets),
207        Some(in_bounds),
208        None,
209        &[],
210        None,
211        None,
212        None,
213        false,
214    );
215    let x = T::load(
216        x_ptr.add_offsets(offsets),
217        Some(in_bounds),
218        None,
219        &[],
220        None,
221        None,
222        None,
223        false,
224    );
225    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
226    let two = T::full(&[BLOCK_SIZE], D::from_f64(2.0));
227    let neg1 = T::full(&[BLOCK_SIZE], D::from_f64(-1.0));
228    let neg2 = T::full(&[BLOCK_SIZE], D::from_f64(-2.0));
229    let sp = T::log(one + T::exp(x));
230    let s2 = one / (one + T::exp(neg2 * sp));
231    let t = two * s2 - one; // tanh(sp)
232    // sigmoid(x) = 1 / (1 + exp(-x))
233    let s = one / (one + T::exp(neg1 * x));
234    // dx = t + x * (1 - t²) * s
235    let dx = dy * (t + x * (one - t * t) * s);
236    T::store(
237        dx_ptr.add_offsets(offsets),
238        dx,
239        Some(in_bounds),
240        &[],
241        None,
242        None,
243    );
244}
245
246pub struct GeluOp<D: Float> {
247    pub forward: GeluForward<D>,
248    pub backward: GeluBackward<D>,
249}
250
251pub struct MishOp<D: Float> {
252    pub forward: MishForward<D>,
253    pub backward: MishBackward<D>,
254}
255
256// ── RuntimeOp for GELU forward ────────────────────────────────────────────────
257
258impl<D: Float + Send + Sync + 'static> teeny_core::model::RuntimeOp for GeluForward<D> {
259    fn n_activation_inputs(&self) -> usize {
260        1
261    }
262
263    fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
264        Vec::new()
265    }
266
267    fn pack_args(
268        &self,
269        inputs: &[(teeny_core::model::RawPtr, &[usize])],
270        _params: &[teeny_core::model::RawPtr],
271        output: teeny_core::model::RawPtr,
272        output_shape: &[usize],
273        _output_row_stride: i32,
274        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
275    ) {
276        let n: usize = output_shape.iter().product();
277        visitor.visit_ptr(inputs[0].0);
278        visitor.visit_ptr(output);
279        visitor.visit_i32(n as i32);
280    }
281
282    fn block(&self) -> [u32; 3] {
283        [self.block_size as u32, 1, 1]
284    }
285
286    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
287        let n: usize = output_shape.iter().product();
288        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
289    }
290
291    #[cfg(feature = "training")]
292    fn has_backward(&self) -> bool {
293        true
294    }
295
296    #[cfg(feature = "training")]
297    fn pack_backward_args(
298        &self,
299        inputs: &[(teeny_core::model::RawPtr, &[usize])],
300        _params: &[teeny_core::model::RawPtr],
301        _output: teeny_core::model::RawPtr,
302        output_shape: &[usize],
303        grad_output: teeny_core::model::RawPtr,
304        _grad_output_row_stride: i32,
305        grad_inputs: &[teeny_core::model::RawPtr],
306        _grad_params: &[teeny_core::model::RawPtr],
307        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
308    ) {
309        let n: usize = output_shape.iter().product();
310        visitor.visit_ptr(grad_output); // dy_ptr
311        visitor.visit_ptr(inputs[0].0); // x_ptr (saved activation)
312        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
313        visitor.visit_i32(n as i32);
314    }
315
316    #[cfg(feature = "training")]
317    fn backward_block(&self) -> [u32; 3] {
318        [self.block_size as u32, 1, 1]
319    }
320
321    #[cfg(feature = "training")]
322    fn backward_grid(&self, _: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
323        let n: usize = output_shape.iter().product();
324        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
325    }
326}
teeny_kernels/nn/activation/gelu.rs

teeny_kernels/nn/activation/
gelu.rs
