teeny_kernels/nn/activation/
relu.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 core::marker::PhantomData;
20use teeny_core::dtype::Num;
21use teeny_macros::kernel;
22use teeny_triton::triton::{
23    types::{AddOffsets, Comparison},
24    *,
25};
26
27#[kernel]
28pub fn relu_forward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
29    x_ptr: T::Pointer<D>,
30    y_ptr: T::Pointer<D>,
31    n_elements: i32,
32) where
33    T::I32Tensor: types::Tensor<i32, 1>,
34    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
35    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
36{
37    let pid = T::program_id(Axis::X);
38    let block_start = pid * BLOCK_SIZE;
39    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
40    let in_bounds = offsets.lt(n_elements);
41
42    let x = T::load(
43        x_ptr.add_offsets(offsets),
44        Some(in_bounds),
45        None,
46        &[],
47        None,
48        None,
49        None,
50        false,
51    );
52    let y = T::zeros_like(x);
53    let relu = T::maximum(x, y);
54
55    // Masked loads in Triton return 0 for masked-off lanes, which gives ReLU.
56    T::store(
57        y_ptr.add_offsets(offsets),
58        relu,
59        Some(in_bounds),
60        &[],
61        None,
62        None,
63    );
64}
65
66#[kernel]
67pub fn relu_backward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
68    dy_ptr: T::Pointer<D>,
69    y_ptr: T::Pointer<D>,
70    dx_ptr: T::Pointer<D>,
71    n_elements: i32,
72) where
73    T::I32Tensor: types::Tensor<i32, 1>,
74    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
75    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
76{
77    let pid = T::program_id(Axis::X);
78    let block_start = pid * BLOCK_SIZE;
79    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
80    let in_bounds = offsets.lt(n_elements);
81
82    let grad_y = T::load(
83        dy_ptr.add_offsets(offsets),
84        Some(in_bounds),
85        None,
86        &[],
87        None,
88        None,
89        None,
90        false,
91    );
92
93    let y = T::load(
94        y_ptr.add_offsets(offsets),
95        Some(in_bounds),
96        None,
97        &[],
98        None,
99        None,
100        None,
101        false,
102    );
103
104    // where(y > 0, grad_y, 0) compiles to a predicated select; avoids an fp mul.
105    let zeros = T::zeros_like(grad_y);
106    let y_gt_zero = T::gt(y, T::zeros_like(y));
107    let grad_x = T::where_(y_gt_zero, grad_y, zeros);
108
109    T::store(
110        dx_ptr.add_offsets(offsets),
111        grad_x,
112        Some(in_bounds),
113        &[],
114        None,
115        None,
116    );
117}
118
119impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for ReluForward<D> {
120    fn n_activation_inputs(&self) -> usize {
121        1
122    }
123
124    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
125        Vec::new()
126    }
127
128    fn pack_args(
129        &self,
130        inputs: &[(teeny_core::model::RawPtr, &[usize])],
131        _params: &[teeny_core::model::RawPtr],
132        output: teeny_core::model::RawPtr,
133        output_shape: &[usize],
134        _output_row_stride: i32,
135        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
136    ) {
137        let n: usize = output_shape.iter().product();
138        visitor.visit_ptr(inputs[0].0);
139        visitor.visit_ptr(output);
140        visitor.visit_i32(n as i32);
141    }
142
143    fn block(&self) -> [u32; 3] {
144        [128, 1, 1]
145    }
146
147    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
148        let n: usize = output_shape.iter().product();
149        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
150    }
151
152    #[cfg(feature = "training")]
153    fn has_backward(&self) -> bool {
154        true
155    }
156
157    // relu_backward(dy_ptr, y_ptr, dx_ptr, n_elements)
158    // dy_ptr = incoming gradient, y_ptr = forward output (activation), dx_ptr = outgoing gradient
159    #[cfg(feature = "training")]
160    fn pack_backward_args(
161        &self,
162        _inputs: &[(teeny_core::model::RawPtr, &[usize])],
163        _params: &[teeny_core::model::RawPtr],
164        output: teeny_core::model::RawPtr,
165        output_shape: &[usize],
166        grad_output: teeny_core::model::RawPtr,
167        _grad_output_row_stride: i32,
168        grad_inputs: &[teeny_core::model::RawPtr],
169        _grad_params: &[teeny_core::model::RawPtr],
170        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
171    ) {
172        let n: usize = output_shape.iter().product();
173        visitor.visit_ptr(grad_output); // dy_ptr
174        visitor.visit_ptr(output); // y_ptr (forward output as activation mask)
175        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
176        visitor.visit_i32(n as i32); // n_elements
177    }
178
179    #[cfg(feature = "training")]
180    fn backward_block(&self) -> [u32; 3] {
181        [128, 1, 1]
182    }
183
184    #[cfg(feature = "training")]
185    fn backward_grid(&self, _input_shapes: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
186        let n: usize = output_shape.iter().product();
187        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
188    }
189}
190
191pub struct ReluOp<'a, T: Num> {
192    pub forward: ReluForward<T>,
193    pub backward: ReluBackward<T>,
194    _marker: PhantomData<&'a ()>,
195}
teeny_kernels/nn/activation/relu.rs

teeny_kernels/nn/activation/
relu.rs
