teeny_kernels/nn/activation/
sigmoid.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// ── Sigmoid ──────────────────────────────────────────────────────────────────
27
28/// Forward: y = 1 / (1 + exp(-x))
29#[kernel(backward = SigmoidBackward)]
30pub fn sigmoid_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
31    x_ptr: T::Pointer<D>,
32    y_ptr: T::Pointer<D>,
33    n_elements: i32,
34) where
35    T::I32Tensor: types::Tensor<i32, 1>,
36    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
37    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
38{
39    let pid = T::program_id(Axis::X);
40    let block_start = pid * BLOCK_SIZE;
41    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
42    let in_bounds = offsets.lt(n_elements);
43
44    let x = T::load(
45        x_ptr.add_offsets(offsets),
46        Some(in_bounds),
47        None,
48        &[],
49        None,
50        None,
51        None,
52        false,
53    );
54    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
55    let neg1 = T::full(&[BLOCK_SIZE], D::from_f64(-1.0));
56    let y = one / (one + T::exp(neg1 * x));
57    T::store(
58        y_ptr.add_offsets(offsets),
59        y,
60        Some(in_bounds),
61        &[],
62        None,
63        None,
64    );
65}
66
67/// Backward: dx = dy * y * (1 - y) = dy * (y - y²)
68#[kernel]
69pub fn sigmoid_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
70    dy_ptr: T::Pointer<D>,
71    y_ptr: T::Pointer<D>,
72    dx_ptr: T::Pointer<D>,
73    n_elements: i32,
74) where
75    T::I32Tensor: types::Tensor<i32, 1>,
76    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
77    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
78{
79    let pid = T::program_id(Axis::X);
80    let block_start = pid * BLOCK_SIZE;
81    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
82    let in_bounds = offsets.lt(n_elements);
83
84    let dy = T::load(
85        dy_ptr.add_offsets(offsets),
86        Some(in_bounds),
87        None,
88        &[],
89        None,
90        None,
91        None,
92        false,
93    );
94    let y = T::load(
95        y_ptr.add_offsets(offsets),
96        Some(in_bounds),
97        None,
98        &[],
99        None,
100        None,
101        None,
102        false,
103    );
104
105    let dx = dy * (y - y * y);
106    T::store(
107        dx_ptr.add_offsets(offsets),
108        dx,
109        Some(in_bounds),
110        &[],
111        None,
112        None,
113    );
114}
115
116// ── SiLU (Swish) ─────────────────────────────────────────────────────────────
117
118/// Forward: y = x * sigmoid(x)
119#[kernel(backward = SiluBackward)]
120pub fn silu_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
121    x_ptr: T::Pointer<D>,
122    y_ptr: T::Pointer<D>,
123    n_elements: i32,
124) where
125    T::I32Tensor: types::Tensor<i32, 1>,
126    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
127    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
128{
129    let pid = T::program_id(Axis::X);
130    let block_start = pid * BLOCK_SIZE;
131    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
132    let in_bounds = offsets.lt(n_elements);
133
134    let x = T::load(
135        x_ptr.add_offsets(offsets),
136        Some(in_bounds),
137        None,
138        &[],
139        None,
140        None,
141        None,
142        false,
143    );
144    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
145    let neg1 = T::full(&[BLOCK_SIZE], D::from_f64(-1.0));
146    let s = one / (one + T::exp(neg1 * x));
147    let y = x * s;
148    T::store(
149        y_ptr.add_offsets(offsets),
150        y,
151        Some(in_bounds),
152        &[],
153        None,
154        None,
155    );
156}
157
158/// Backward: dx = dy * (sigmoid(x) + y * (1 - sigmoid(x)))
159///         = dy * (s + y - y*s)   where s = sigmoid(x)
160#[kernel]
161pub fn silu_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
162    dy_ptr: T::Pointer<D>,
163    x_ptr: T::Pointer<D>,
164    dx_ptr: T::Pointer<D>,
165    n_elements: i32,
166) where
167    T::I32Tensor: types::Tensor<i32, 1>,
168    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
169    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
170{
171    let pid = T::program_id(Axis::X);
172    let block_start = pid * BLOCK_SIZE;
173    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
174    let in_bounds = offsets.lt(n_elements);
175
176    let dy = T::load(
177        dy_ptr.add_offsets(offsets),
178        Some(in_bounds),
179        None,
180        &[],
181        None,
182        None,
183        None,
184        false,
185    );
186    let x = T::load(
187        x_ptr.add_offsets(offsets),
188        Some(in_bounds),
189        None,
190        &[],
191        None,
192        None,
193        None,
194        false,
195    );
196    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
197    let neg1 = T::full(&[BLOCK_SIZE], D::from_f64(-1.0));
198    let s = one / (one + T::exp(neg1 * x));
199    let y = x * s;
200    // d(silu)/dx = s + x*s*(1-s) = s + y - y*s
201    let dx = dy * (s + y - y * s);
202    T::store(
203        dx_ptr.add_offsets(offsets),
204        dx,
205        Some(in_bounds),
206        &[],
207        None,
208        None,
209    );
210}
211
212// ── LogSigmoid ────────────────────────────────────────────────────────────────
213
214/// Forward: y = log(sigmoid(x)) = -log(1 + exp(-x))
215#[kernel(backward = LogsigmoidBackward)]
216pub fn logsigmoid_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
217    x_ptr: T::Pointer<D>,
218    y_ptr: T::Pointer<D>,
219    n_elements: i32,
220) where
221    T::I32Tensor: types::Tensor<i32, 1>,
222    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
223    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
224{
225    let pid = T::program_id(Axis::X);
226    let block_start = pid * BLOCK_SIZE;
227    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
228    let in_bounds = offsets.lt(n_elements);
229
230    let x = T::load(
231        x_ptr.add_offsets(offsets),
232        Some(in_bounds),
233        None,
234        &[],
235        None,
236        None,
237        None,
238        false,
239    );
240    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
241    let neg1 = T::full(&[BLOCK_SIZE], D::from_f64(-1.0));
242    // -log(1 + exp(-x)) = log(1/(1+exp(-x))) = log(sigmoid(x))
243    // But we want to avoid negating the result: use (neg1 * log(1 + exp(neg1*x)))
244    // Actually: y = neg1 * log(one + T::exp(neg1 * x))
245    // But neg1 * log(...) would require negating a tensor result.
246    // Use subtraction: y = T::zeros_like(x) - T::log(one + T::exp(neg1 * x))
247    let zeros = T::zeros_like(x);
248    let y = zeros - T::log(one + T::exp(neg1 * x));
249    T::store(
250        y_ptr.add_offsets(offsets),
251        y,
252        Some(in_bounds),
253        &[],
254        None,
255        None,
256    );
257}
258
259/// Backward: dx = dy * sigmoid(-x) = dy / (1 + exp(x))
260#[kernel]
261pub fn logsigmoid_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
262    dy_ptr: T::Pointer<D>,
263    x_ptr: T::Pointer<D>,
264    dx_ptr: T::Pointer<D>,
265    n_elements: i32,
266) where
267    T::I32Tensor: types::Tensor<i32, 1>,
268    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
269    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
270{
271    let pid = T::program_id(Axis::X);
272    let block_start = pid * BLOCK_SIZE;
273    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
274    let in_bounds = offsets.lt(n_elements);
275
276    let dy = T::load(
277        dy_ptr.add_offsets(offsets),
278        Some(in_bounds),
279        None,
280        &[],
281        None,
282        None,
283        None,
284        false,
285    );
286    let x = T::load(
287        x_ptr.add_offsets(offsets),
288        Some(in_bounds),
289        None,
290        &[],
291        None,
292        None,
293        None,
294        false,
295    );
296    let one = T::full(&[BLOCK_SIZE], D::from_f64(1.0));
297    // sigmoid(-x) = 1 / (1 + exp(x))
298    let dx = dy / (one + T::exp(x));
299    T::store(
300        dx_ptr.add_offsets(offsets),
301        dx,
302        Some(in_bounds),
303        &[],
304        None,
305        None,
306    );
307}
308
309pub struct SigmoidOp<D: Float> {
310    pub forward: SigmoidForward<D>,
311    pub backward: SigmoidBackward<D>,
312}
313
314pub struct SiluOp<D: Float> {
315    pub forward: SiluForward<D>,
316    pub backward: SiluBackward<D>,
317}
318
319pub struct LogsigmoidOp<D: Float> {
320    pub forward: LogsigmoidForward<D>,
321    pub backward: LogsigmoidBackward<D>,
322}
323
324// ── RuntimeOp for Sigmoid forward ────────────────────────────────────────────
325
326impl<D: Float + Send + Sync + 'static> teeny_core::model::RuntimeOp for SigmoidForward<D> {
327    fn n_activation_inputs(&self) -> usize {
328        1
329    }
330
331    fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
332        Vec::new()
333    }
334
335    fn pack_args(
336        &self,
337        inputs: &[(teeny_core::model::RawPtr, &[usize])],
338        _params: &[teeny_core::model::RawPtr],
339        output: teeny_core::model::RawPtr,
340        output_shape: &[usize],
341        _output_row_stride: i32,
342        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
343    ) {
344        let n: usize = output_shape.iter().product();
345        visitor.visit_ptr(inputs[0].0);
346        visitor.visit_ptr(output);
347        visitor.visit_i32(n as i32);
348    }
349
350    fn block(&self) -> [u32; 3] {
351        [self.block_size as u32, 1, 1]
352    }
353
354    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
355        let n: usize = output_shape.iter().product();
356        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
357    }
358
359    #[cfg(feature = "training")]
360    fn has_backward(&self) -> bool {
361        true
362    }
363
364    #[cfg(feature = "training")]
365    fn pack_backward_args(
366        &self,
367        _inputs: &[(teeny_core::model::RawPtr, &[usize])],
368        _params: &[teeny_core::model::RawPtr],
369        output: teeny_core::model::RawPtr,
370        output_shape: &[usize],
371        grad_output: teeny_core::model::RawPtr,
372        _grad_output_row_stride: i32,
373        grad_inputs: &[teeny_core::model::RawPtr],
374        _grad_params: &[teeny_core::model::RawPtr],
375        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
376    ) {
377        let n: usize = output_shape.iter().product();
378        visitor.visit_ptr(grad_output); // dy_ptr
379        visitor.visit_ptr(output); // y_ptr (saved output, not x)
380        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
381        visitor.visit_i32(n as i32);
382    }
383
384    #[cfg(feature = "training")]
385    fn backward_block(&self) -> [u32; 3] {
386        [self.block_size as u32, 1, 1]
387    }
388
389    #[cfg(feature = "training")]
390    fn backward_grid(&self, _: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
391        let n: usize = output_shape.iter().product();
392        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
393    }
394}
395
396// ── RuntimeOp for SiLU forward ────────────────────────────────────────────────
397
398impl<D: Float + Send + Sync + 'static> teeny_core::model::RuntimeOp for SiluForward<D> {
399    fn n_activation_inputs(&self) -> usize {
400        1
401    }
402
403    fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
404        Vec::new()
405    }
406
407    fn pack_args(
408        &self,
409        inputs: &[(teeny_core::model::RawPtr, &[usize])],
410        _params: &[teeny_core::model::RawPtr],
411        output: teeny_core::model::RawPtr,
412        output_shape: &[usize],
413        _output_row_stride: i32,
414        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
415    ) {
416        let n: usize = output_shape.iter().product();
417        visitor.visit_ptr(inputs[0].0);
418        visitor.visit_ptr(output);
419        visitor.visit_i32(n as i32);
420    }
421
422    fn block(&self) -> [u32; 3] {
423        [self.block_size as u32, 1, 1]
424    }
425
426    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
427        let n: usize = output_shape.iter().product();
428        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
429    }
430
431    #[cfg(feature = "training")]
432    fn has_backward(&self) -> bool {
433        true
434    }
435
436    #[cfg(feature = "training")]
437    fn pack_backward_args(
438        &self,
439        inputs: &[(teeny_core::model::RawPtr, &[usize])],
440        _params: &[teeny_core::model::RawPtr],
441        _output: teeny_core::model::RawPtr,
442        output_shape: &[usize],
443        grad_output: teeny_core::model::RawPtr,
444        _grad_output_row_stride: i32,
445        grad_inputs: &[teeny_core::model::RawPtr],
446        _grad_params: &[teeny_core::model::RawPtr],
447        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
448    ) {
449        let n: usize = output_shape.iter().product();
450        visitor.visit_ptr(grad_output); // dy_ptr
451        visitor.visit_ptr(inputs[0].0); // x_ptr (saved activation)
452        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
453        visitor.visit_i32(n as i32);
454    }
455
456    #[cfg(feature = "training")]
457    fn backward_block(&self) -> [u32; 3] {
458        [self.block_size as u32, 1, 1]
459    }
460
461    #[cfg(feature = "training")]
462    fn backward_grid(&self, _: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
463        let n: usize = output_shape.iter().product();
464        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
465    }
466}
teeny_kernels/nn/activation/sigmoid.rs

teeny_kernels/nn/activation/
sigmoid.rs
