teeny_kernels/nn/conv/
conv1d.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::ops::BitAnd;
20
21use teeny_core::dtype::Num;
22use teeny_macros::kernel;
23use teeny_triton::triton::{
24    types::{AddOffsets, Comparison, Tensor},
25    *,
26};
27
28/// 1-D convolution forward pass.
29///
30/// Grid: one CTA per (b, c_out, ol-tile):
31///   `pid = (b * C_OUT + c_out) * num_ol_tiles + ol_tile`
32///
33/// Each CTA computes a BLOCK_OL-wide strip of output positions by iterating
34/// over all `C_IN * KL` combinations.
35///
36/// Zero-padding of `PAD` elements is applied on each side of the input.
37/// `OL = (L + 2*PAD - KL) / STRIDE + 1`.
38#[kernel]
39pub fn conv1d_forward<
40    T: Triton,
41    D: Num,
42    const KL: i32,
43    const STRIDE: i32,
44    const PAD: i32,
45    const BLOCK_OL: i32,
46>(
47    x_ptr: T::Pointer<D>,
48    w_ptr: T::Pointer<D>,
49    y_ptr: T::Pointer<D>,
50    _B: i32,
51    C_IN: i32,
52    C_OUT: i32,
53    L: i32,
54    OL: i32,
55) where
56    T::I32Tensor: Tensor<i32, 1>,
57    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
58    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
59    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
60{
61    let pid = T::program_id(Axis::X);
62    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
63
64    let ol_tile = pid % num_ol_tiles;
65    let bc = pid / num_ol_tiles;
66    let c_out = bc % C_OUT;
67    let b = bc / C_OUT;
68
69    let ol_start = ol_tile * BLOCK_OL;
70    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
71    let ol_mask = ol_range.lt(OL);
72
73    let out_bc_base = (b * C_OUT + c_out) * OL;
74
75    let mut acc = T::zeros::<D>(&[BLOCK_OL]);
76
77    let loop_bound = C_IN * KL;
78    for idx in 0..loop_bound {
79        let kl = idx % KL;
80        let c_in = idx / KL;
81
82        let il_range = ol_range * STRIDE + kl - PAD;
83        let in_bounds = il_range.ge(0) & il_range.lt(L);
84        let load_mask = ol_mask & in_bounds;
85
86        let x_offsets = il_range + (b * C_IN + c_in) * L;
87        let x_tile = T::load(
88            x_ptr.add_offsets(x_offsets),
89            Some(load_mask),
90            Some(T::zeros::<D>(&[BLOCK_OL])),
91            &[],
92            None,
93            None,
94            None,
95            false,
96        );
97
98        let w_idx = (c_out * C_IN + c_in) * KL + kl;
99        let w_off = T::arange(0, 1) + w_idx;
100        let w_1 = T::load(
101            w_ptr.add_offsets(w_off),
102            None,
103            None,
104            &[],
105            None,
106            None,
107            None,
108            false,
109        );
110        let w_tile = T::broadcast_to(w_1, &[BLOCK_OL]);
111
112        acc = acc + x_tile * w_tile;
113    }
114
115    let out_offsets = ol_range + out_bc_base;
116    T::store(
117        y_ptr.add_offsets(out_offsets),
118        acc,
119        Some(ol_mask),
120        &[],
121        None,
122        None,
123    );
124}
125
126/// 1-D convolution backward pass — gradient with respect to input (`dx`).
127///
128/// Grid: `pid = (b * C_OUT + c_out) * num_ol_tiles + ol_tile`
129///
130/// Scatters gradient back via `atomic_add` to handle overlapping receptive fields.
131/// Padding positions (those that correspond to out-of-bounds input locations)
132/// are skipped.
133#[kernel]
134pub fn conv1d_backward_dx<
135    T: Triton,
136    D: Num,
137    const KL: i32,
138    const STRIDE: i32,
139    const PAD: i32,
140    const BLOCK_OL: i32,
141>(
142    dy_ptr: T::Pointer<D>,
143    w_ptr: T::Pointer<D>,
144    dx_ptr: T::Pointer<D>,
145    _B: i32,
146    C_IN: i32,
147    C_OUT: i32,
148    L: i32,
149    OL: i32,
150) where
151    T::I32Tensor: Tensor<i32, 1>,
152    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
153    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
154    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
155{
156    let pid = T::program_id(Axis::X);
157    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
158
159    let ol_tile = pid % num_ol_tiles;
160    let bc = pid / num_ol_tiles;
161    let c_out = bc % C_OUT;
162    let b = bc / C_OUT;
163
164    let ol_start = ol_tile * BLOCK_OL;
165    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
166    let ol_mask = ol_range.lt(OL);
167
168    let dy_offsets = ol_range + (b * C_OUT + c_out) * OL;
169    let dy_tile = T::load(
170        dy_ptr.add_offsets(dy_offsets),
171        Some(ol_mask),
172        Some(T::zeros::<D>(&[BLOCK_OL])),
173        &[],
174        None,
175        None,
176        None,
177        false,
178    );
179
180    let loop_bound = C_IN * KL;
181    for idx in 0..loop_bound {
182        let kl = idx % KL;
183        let c_in = idx / KL;
184
185        let w_idx = (c_out * C_IN + c_in) * KL + kl;
186        let w_off = T::arange(0, 1) + w_idx;
187        let w_1 = T::load(
188            w_ptr.add_offsets(w_off),
189            None,
190            None,
191            &[],
192            None,
193            None,
194            None,
195            false,
196        );
197        let w_tile = T::broadcast_to(w_1, &[BLOCK_OL]);
198
199        let grad_tile = dy_tile * w_tile;
200
201        let il_range = ol_range * STRIDE + kl - PAD;
202        let in_bounds = il_range.ge(0) & il_range.lt(L);
203        let dx_offsets = il_range + (b * C_IN + c_in) * L;
204        T::atomic_add(
205            dx_ptr.add_offsets(dx_offsets),
206            grad_tile,
207            Some(ol_mask & in_bounds),
208            None,
209            None,
210        );
211    }
212}
213
214/// 1-D convolution backward pass — gradient with respect to weights (`dw`).
215///
216/// Grid: `pid = (b * C_OUT + c_out) * num_ol_tiles + ol_tile`
217///
218/// `dw` must be zero-initialised before launch.
219#[kernel]
220pub fn conv1d_backward_dw<
221    T: Triton,
222    D: Num,
223    const KL: i32,
224    const STRIDE: i32,
225    const PAD: i32,
226    const BLOCK_OL: i32,
227>(
228    dy_ptr: T::Pointer<D>,
229    x_ptr: T::Pointer<D>,
230    dw_ptr: T::Pointer<D>,
231    _B: i32,
232    C_IN: i32,
233    C_OUT: i32,
234    L: i32,
235    OL: i32,
236) where
237    T::I32Tensor: Tensor<i32, 1>,
238    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
239    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
240    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
241{
242    let pid = T::program_id(Axis::X);
243    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
244
245    let ol_tile = pid % num_ol_tiles;
246    let bc = pid / num_ol_tiles;
247    let c_out = bc % C_OUT;
248    let b = bc / C_OUT;
249
250    let ol_start = ol_tile * BLOCK_OL;
251    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
252    let ol_mask = ol_range.lt(OL);
253
254    let dy_offsets = ol_range + (b * C_OUT + c_out) * OL;
255    let dy_tile = T::load(
256        dy_ptr.add_offsets(dy_offsets),
257        Some(ol_mask),
258        Some(T::zeros::<D>(&[BLOCK_OL])),
259        &[],
260        None,
261        None,
262        None,
263        false,
264    );
265
266    let loop_bound = C_IN * KL;
267    for idx in 0..loop_bound {
268        let kl = idx % KL;
269        let c_in = idx / KL;
270
271        let il_range = ol_range * STRIDE + kl - PAD;
272        let in_bounds = il_range.ge(0) & il_range.lt(L);
273        let load_mask = ol_mask & in_bounds;
274
275        let x_offsets = il_range + (b * C_IN + c_in) * L;
276        let x_tile = T::load(
277            x_ptr.add_offsets(x_offsets),
278            Some(load_mask),
279            Some(T::zeros::<D>(&[BLOCK_OL])),
280            &[],
281            None,
282            None,
283            None,
284            false,
285        );
286
287        let partial = T::sum(dy_tile * x_tile, Some(0), false);
288        let partial_1 = T::expand_dims(partial, 0);
289
290        let w_idx = (c_out * C_IN + c_in) * KL + kl;
291        let dw_off = T::arange(0, 1) + w_idx;
292        T::atomic_add(dw_ptr.add_offsets(dw_off), partial_1, None, None, None);
293    }
294}
295
296impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for Conv1dForward<D> {
297    fn n_activation_inputs(&self) -> usize {
298        1
299    }
300
301    fn param_shapes(&self, input_shapes: &[&[usize]], output_shape: &[usize]) -> Vec<Vec<usize>> {
302        // input_shapes[0] = [B, C_IN, L], output_shape = [B, C_OUT, OL]
303        let c_in = input_shapes[0][1];
304        let c_out = output_shape[1];
305        vec![vec![c_out, c_in, self.kl as usize]]
306    }
307
308    fn pack_args(
309        &self,
310        inputs: &[(teeny_core::model::RawPtr, &[usize])],
311        params: &[teeny_core::model::RawPtr],
312        output: teeny_core::model::RawPtr,
313        output_shape: &[usize],
314        _output_row_stride: i32,
315        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
316    ) {
317        // kernel args: x_ptr, w_ptr, y_ptr, B, C_IN, C_OUT, L, OL
318        let input_shape = inputs[0].1;
319        visitor.visit_ptr(inputs[0].0);
320        visitor.visit_ptr(params[0]);
321        visitor.visit_ptr(output);
322        visitor.visit_i32(input_shape[0] as i32); // B
323        visitor.visit_i32(input_shape[1] as i32); // C_IN
324        visitor.visit_i32(output_shape[1] as i32); // C_OUT
325        visitor.visit_i32(input_shape[2] as i32); // L
326        visitor.visit_i32(output_shape[2] as i32); // OL
327    }
328
329    fn block(&self) -> [u32; 3] {
330        [128, 1, 1]
331    }
332
333    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
334        let num_ol_tiles = output_shape[2].div_ceil(self.block_ol as usize);
335        [
336            (output_shape[0] * output_shape[1] * num_ol_tiles) as u32,
337            1,
338            1,
339        ]
340    }
341}
342
343pub struct Conv1dOp<'a, T: Num> {
344    pub forward: Conv1dForward<T>,
345    pub backward_dx: Conv1dBackwardDx<T>,
346    pub backward_dw: Conv1dBackwardDw<T>,
347    _marker: core::marker::PhantomData<&'a ()>,
348}
teeny_kernels/nn/conv/conv1d.rs

teeny_kernels/nn/conv/
conv1d.rs
