teeny_kernels/nn/tensor/
elemwise_add.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::Num;
20use teeny_macros::kernel;
21use teeny_triton::triton::{
22    types::{AddOffsets, Comparison},
23    *,
24};
25
26// ── Forward: out[i] = a[i] + b[i] ────────────────────────────────────────────
27
28#[kernel]
29pub fn elemwise_add_forward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
30    a_ptr: T::Pointer<D>,
31    b_ptr: T::Pointer<D>,
32    out_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 a = T::load(
45        a_ptr.add_offsets(offsets),
46        Some(in_bounds),
47        None,
48        &[],
49        None,
50        None,
51        None,
52        false,
53    );
54    let b = T::load(
55        b_ptr.add_offsets(offsets),
56        Some(in_bounds),
57        None,
58        &[],
59        None,
60        None,
61        None,
62        false,
63    );
64    T::store(
65        out_ptr.add_offsets(offsets),
66        a + b,
67        Some(in_bounds),
68        &[],
69        None,
70        None,
71    );
72}
73
74// ── Backward: grad_a[i] = dy[i],  grad_b[i] = dy[i] ─────────────────────────
75//
76// Add is the fan-out of the gradient: the upstream gradient flows unchanged
77// to both inputs.
78
79#[kernel]
80pub fn elemwise_add_backward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
81    dy_ptr: T::Pointer<D>,
82    grad_a_ptr: T::Pointer<D>,
83    grad_b_ptr: T::Pointer<D>,
84    n_elements: i32,
85) where
86    T::I32Tensor: types::Tensor<i32, 1>,
87    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
88    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
89{
90    let pid = T::program_id(Axis::X);
91    let block_start = pid * BLOCK_SIZE;
92    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
93    let in_bounds = offsets.lt(n_elements);
94
95    let dy = T::load(
96        dy_ptr.add_offsets(offsets),
97        Some(in_bounds),
98        None,
99        &[],
100        None,
101        None,
102        None,
103        false,
104    );
105    T::store(
106        grad_a_ptr.add_offsets(offsets),
107        dy,
108        Some(in_bounds),
109        &[],
110        None,
111        None,
112    );
113    T::store(
114        grad_b_ptr.add_offsets(offsets),
115        dy,
116        Some(in_bounds),
117        &[],
118        None,
119        None,
120    );
121}
122
123// ── RuntimeOp ─────────────────────────────────────────────────────────────────
124
125impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for ElemwiseAddForward<D> {
126    fn n_activation_inputs(&self) -> usize {
127        2
128    }
129
130    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
131        Vec::new()
132    }
133
134    fn pack_args(
135        &self,
136        inputs: &[(teeny_core::model::RawPtr, &[usize])],
137        _params: &[teeny_core::model::RawPtr],
138        output: teeny_core::model::RawPtr,
139        output_shape: &[usize],
140        _output_row_stride: i32,
141        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
142    ) {
143        let n: usize = output_shape.iter().product();
144        visitor.visit_ptr(inputs[0].0); // a_ptr
145        visitor.visit_ptr(inputs[1].0); // b_ptr
146        visitor.visit_ptr(output); // out_ptr
147        visitor.visit_i32(n as i32); // n_elements
148    }
149
150    fn block(&self) -> [u32; 3] {
151        [self.block_size as u32, 1, 1]
152    }
153
154    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
155        let n: usize = output_shape.iter().product();
156        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
157    }
158
159    #[cfg(feature = "training")]
160    fn has_backward(&self) -> bool {
161        true
162    }
163
164    #[cfg(feature = "training")]
165    fn pack_backward_args(
166        &self,
167        _inputs: &[(teeny_core::model::RawPtr, &[usize])],
168        _params: &[teeny_core::model::RawPtr],
169        _output: teeny_core::model::RawPtr,
170        output_shape: &[usize],
171        grad_output: teeny_core::model::RawPtr,
172        _grad_output_row_stride: i32,
173        grad_inputs: &[teeny_core::model::RawPtr],
174        _grad_params: &[teeny_core::model::RawPtr],
175        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
176    ) {
177        let n: usize = output_shape.iter().product();
178        visitor.visit_ptr(grad_output); // dy_ptr
179        visitor.visit_ptr(grad_inputs[0]); // grad_a_ptr
180        visitor.visit_ptr(grad_inputs[1]); // grad_b_ptr
181        visitor.visit_i32(n as i32); // n_elements
182    }
183
184    #[cfg(feature = "training")]
185    fn backward_block(&self) -> [u32; 3] {
186        [self.block_size as u32, 1, 1]
187    }
188
189    #[cfg(feature = "training")]
190    fn backward_grid(&self, _input_shapes: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
191        let n: usize = output_shape.iter().product();
192        [n.div_ceil(self.block_size as usize) as u32, 1, 1]
193    }
194}
teeny_kernels/nn/tensor/elemwise_add.rs

teeny_kernels/nn/tensor/
elemwise_add.rs
