teeny_kernels/nn/conv/
conv2d.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/// 2-D convolution forward pass (supports grouped and depthwise conv).
29///
30/// Grid: one CTA per (b, c_out, oh, ow-tile):
31///   `pid = ((b * C_OUT + c_out) * OH + oh) * num_ow_tiles + ow_tile`
32///
33/// Each CTA computes a BLOCK_OW-wide strip of output columns by iterating over
34/// `(C_IN/G) * KH * KW` combinations for its assigned group.  The weight for
35/// each `(c_in_local, kh, kw)` is loaded as a [1] tensor and broadcast to
36/// `[BLOCK_OW]`.
37///
38/// `G` is the number of groups (1 = standard conv, G = C_IN = C_OUT for depthwise).
39/// Weight layout: `[C_OUT, C_IN/G, KH, KW]`.
40///
41/// Zero-padding of `PAD_H` / `PAD_W` elements is applied on each spatial side.
42/// `OH = (H + 2*PAD_H - KH) / STRIDE_H + 1`, `OW = (W + 2*PAD_W - KW) / STRIDE_W + 1`.
43#[kernel]
44pub fn conv2d_forward<
45    T: Triton,
46    D: Num,
47    const KH: i32,
48    const KW: i32,
49    const STRIDE_H: i32,
50    const STRIDE_W: i32,
51    const PAD_H: i32,
52    const PAD_W: i32,
53    const G: i32,
54    const BLOCK_OW: i32,
55>(
56    x_ptr: T::Pointer<D>,
57    w_ptr: T::Pointer<D>,
58    y_ptr: T::Pointer<D>,
59    _B: i32,
60    C_IN: i32,
61    C_OUT: i32,
62    H: i32,
63    W: i32,
64    OH: i32,
65    OW: i32,
66) where
67    T::I32Tensor: Tensor<i32, 1>,
68    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
69    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
70    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
71{
72    let pid = T::program_id(Axis::X);
73    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
74
75    // Decode flat pid → (b, c_out, oh, ow_tile).
76    let ow_tile = pid % num_ow_tiles;
77    let bco = pid / num_ow_tiles;
78    let oh = bco % OH;
79    let bc = bco / OH;
80    let c_out = bc % C_OUT;
81    let b = bc / C_OUT;
82
83    let ow_start = ow_tile * BLOCK_OW;
84    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
85    let ow_mask = ow_range.lt(OW);
86
87    let out_bc_base = (b * C_OUT + c_out) * OH * OW;
88
89    // Group index for this output channel and its input-channel window.
90    let c_in_per_group = C_IN / G;
91    let g_idx = c_out / (C_OUT / G);
92    let c_in_start = g_idx * c_in_per_group;
93
94    let mut acc = T::zeros::<D>(&[BLOCK_OW]);
95
96    // Flat loop over (C_IN/G) * KH * KW combinations for this group.
97    let loop_bound = c_in_per_group * KH * KW;
98    for idx in 0..loop_bound {
99        let kw = idx % KW;
100        let kh_cin = idx / KW;
101        let kh = kh_cin % KH;
102        let c_in_local = kh_cin / KH; // index within the group
103        let c_in = c_in_start + c_in_local; // absolute input channel
104
105        // Compute padded input coordinates; OOB height rows contribute zero via mask.
106        let ih = oh * STRIDE_H + kh - PAD_H;
107        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
108
109        // `ow_range * 0` is the only way to splat scalar ih into an I32Tensor.
110        // A scalar `if`/`continue` here triggers a compiler phi-node bug.
111        #[allow(clippy::erasing_op)]
112        let ih_t = ow_range * 0 + ih;
113        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
114        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
115        let load_mask = ow_mask & h_in_bounds & w_in_bounds;
116
117        let x_offsets = iw_range + ((b * C_IN + c_in) * H * W + ih * W);
118        let x_tile = T::load(
119            x_ptr.add_offsets(x_offsets),
120            Some(load_mask),
121            Some(T::zeros::<D>(&[BLOCK_OW])),
122            &[],
123            None,
124            None,
125            None,
126            false,
127        );
128
129        // Weight layout [C_OUT, C_IN/G, KH, KW]: load scalar and broadcast.
130        let w_idx = ((c_out * c_in_per_group + c_in_local) * KH + kh) * KW + kw;
131        let w_off = T::arange(0, 1) + w_idx;
132        let w_1 = T::load(
133            w_ptr.add_offsets(w_off),
134            None,
135            None,
136            &[],
137            None,
138            None,
139            None,
140            false,
141        );
142        let w_tile = T::broadcast_to(w_1, &[BLOCK_OW]);
143
144        acc = acc + x_tile * w_tile;
145    }
146
147    let out_offsets = ow_range + (out_bc_base + oh * OW);
148    T::store(
149        y_ptr.add_offsets(out_offsets),
150        acc,
151        Some(ow_mask),
152        &[],
153        None,
154        None,
155    );
156}
157
158/// 2-D convolution backward pass — gradient with respect to input (`dx`).
159///
160/// Uses the same grid as the forward pass (over output positions) and scatters
161/// the gradient back to the input via `atomic_add`, which correctly handles
162/// overlapping receptive fields when `STRIDE < kernel size`.
163///
164/// `G` — number of groups (1 = standard conv, G = C_IN = C_OUT for depthwise).
165/// Weight layout: `[C_OUT, C_IN/G, KH, KW]`.
166///
167/// Grid: `pid = ((b * C_OUT + c_out) * OH + oh) * num_ow_tiles + ow_tile`
168///
169/// Padding positions (those that correspond to out-of-bounds input locations)
170/// are skipped.
171#[kernel]
172pub fn conv2d_backward_dx<
173    T: Triton,
174    D: Num,
175    const KH: i32,
176    const KW: i32,
177    const STRIDE_H: i32,
178    const STRIDE_W: i32,
179    const PAD_H: i32,
180    const PAD_W: i32,
181    const G: i32,
182    const BLOCK_OW: i32,
183>(
184    dy_ptr: T::Pointer<D>,
185    w_ptr: T::Pointer<D>,
186    dx_ptr: T::Pointer<D>,
187    _B: i32,
188    C_IN: i32,
189    C_OUT: i32,
190    H: i32,
191    W: i32,
192    OH: i32,
193    OW: i32,
194) where
195    T::I32Tensor: Tensor<i32, 1>,
196    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
197    T::BoolTensor: BitAnd<Output = 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 num_ow_tiles = T::cdiv(OW, BLOCK_OW);
202
203    let ow_tile = pid % num_ow_tiles;
204    let bco = pid / num_ow_tiles;
205    let oh = bco % OH;
206    let bc = bco / OH;
207    let c_out = bc % C_OUT;
208    let b = bc / C_OUT;
209
210    let ow_start = ow_tile * BLOCK_OW;
211    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
212    let ow_mask = ow_range.lt(OW);
213
214    // Load the upstream gradient tile for this output row.
215    let dy_offsets = ow_range + ((b * C_OUT + c_out) * OH * OW + oh * OW);
216    let dy_tile = T::load(
217        dy_ptr.add_offsets(dy_offsets),
218        Some(ow_mask),
219        Some(T::zeros::<D>(&[BLOCK_OW])),
220        &[],
221        None,
222        None,
223        None,
224        false,
225    );
226
227    // Group index for this output channel and its input-channel window.
228    let c_in_per_group = C_IN / G;
229    let g_idx = c_out / (C_OUT / G);
230    let c_in_start = g_idx * c_in_per_group;
231
232    // Scatter gradient to input via flat loop over (C_IN/G) * KH * KW.
233    let loop_bound = c_in_per_group * KH * KW;
234    for idx in 0..loop_bound {
235        let kw = idx % KW;
236        let kh_cin = idx / KW;
237        let kh = kh_cin % KH;
238        let c_in_local = kh_cin / KH;
239        let c_in = c_in_start + c_in_local;
240
241        let w_idx = ((c_out * c_in_per_group + c_in_local) * KH + kh) * KW + kw;
242        let w_off = T::arange(0, 1) + w_idx;
243        let w_1 = T::load(
244            w_ptr.add_offsets(w_off),
245            None,
246            None,
247            &[],
248            None,
249            None,
250            None,
251            false,
252        );
253        let w_tile = T::broadcast_to(w_1, &[BLOCK_OW]);
254
255        let grad_tile = dy_tile * w_tile;
256
257        let ih = oh * STRIDE_H + kh - PAD_H;
258        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
259
260        #[allow(clippy::erasing_op)]
261        let ih_t = ow_range * 0 + ih;
262        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
263        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
264
265        let dx_offsets = iw_range + ((b * C_IN + c_in) * H * W + ih * W);
266        T::atomic_add(
267            dx_ptr.add_offsets(dx_offsets),
268            grad_tile,
269            Some(ow_mask & h_in_bounds & w_in_bounds),
270            None,
271            None,
272        );
273    }
274}
275
276/// 2-D convolution backward pass — gradient with respect to weights (`dw`).
277///
278/// Uses the same grid as the forward pass.  Each CTA at `(b, c_out, oh, ow_tile)`
279/// accumulates partial sums into `dw` via `atomic_add` (one per `(c_in_local, kh, kw)`
280/// combination within the group).
281///
282/// `G` — number of groups. Weight layout: `[C_OUT, C_IN/G, KH, KW]`.
283/// Grid: `pid = ((b * C_OUT + c_out) * OH + oh) * num_ow_tiles + ow_tile`
284///
285/// `dw` must be zero-initialised before launch.
286#[kernel]
287pub fn conv2d_backward_dw<
288    T: Triton,
289    D: Num,
290    const KH: i32,
291    const KW: i32,
292    const STRIDE_H: i32,
293    const STRIDE_W: i32,
294    const PAD_H: i32,
295    const PAD_W: i32,
296    const G: i32,
297    const BLOCK_OW: i32,
298>(
299    dy_ptr: T::Pointer<D>,
300    x_ptr: T::Pointer<D>,
301    dw_ptr: T::Pointer<D>,
302    _B: i32,
303    C_IN: i32,
304    C_OUT: i32,
305    H: i32,
306    W: i32,
307    OH: i32,
308    OW: i32,
309) where
310    T::I32Tensor: Tensor<i32, 1>,
311    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
312    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
313    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
314{
315    let pid = T::program_id(Axis::X);
316    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
317
318    let ow_tile = pid % num_ow_tiles;
319    let bco = pid / num_ow_tiles;
320    let oh = bco % OH;
321    let bc = bco / OH;
322    let c_out = bc % C_OUT;
323    let b = bc / C_OUT;
324
325    let ow_start = ow_tile * BLOCK_OW;
326    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
327    let ow_mask = ow_range.lt(OW);
328
329    // Load dy tile for this output row.
330    let dy_offsets = ow_range + ((b * C_OUT + c_out) * OH * OW + oh * OW);
331    let dy_tile = T::load(
332        dy_ptr.add_offsets(dy_offsets),
333        Some(ow_mask),
334        Some(T::zeros::<D>(&[BLOCK_OW])),
335        &[],
336        None,
337        None,
338        None,
339        false,
340    );
341
342    let ih_base = oh * STRIDE_H;
343
344    // Group index for this output channel and its input-channel window.
345    let c_in_per_group = C_IN / G;
346    let g_idx = c_out / (C_OUT / G);
347    let c_in_start = g_idx * c_in_per_group;
348
349    // Flat loop over (C_IN/G) * KH * KW to compute partial weight gradients.
350    let loop_bound = c_in_per_group * KH * KW;
351    for idx in 0..loop_bound {
352        let kw = idx % KW;
353        let kh_cin = idx / KW;
354        let kh = kh_cin % KH;
355        let c_in_local = kh_cin / KH;
356        let c_in = c_in_start + c_in_local;
357
358        let ih = ih_base + kh - PAD_H;
359        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
360
361        #[allow(clippy::erasing_op)]
362        let ih_t = ow_range * 0 + ih;
363        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
364        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
365        let load_mask = ow_mask & h_in_bounds & w_in_bounds;
366
367        let x_offsets = iw_range + ((b * C_IN + c_in) * H * W + ih * W);
368        let x_tile = T::load(
369            x_ptr.add_offsets(x_offsets),
370            Some(load_mask),
371            Some(T::zeros::<D>(&[BLOCK_OW])),
372            &[],
373            None,
374            None,
375            None,
376            false,
377        );
378
379        // Partial sum for this weight element: scalar dot product over the ow tile.
380        let partial = T::sum(dy_tile * x_tile, Some(0), false);
381        let partial_1 = T::expand_dims(partial, 0);
382
383        let w_idx = ((c_out * c_in_per_group + c_in_local) * KH + kh) * KW + kw;
384        let dw_off = T::arange(0, 1) + w_idx;
385        T::atomic_add(dw_ptr.add_offsets(dw_off), partial_1, None, None, None);
386    }
387}
388
389/// 2-D convolution combined backward pass — computes both `dx` and `dw` in one
390/// kernel launch, using the same grid as the forward pass.
391///
392/// Grid: `pid = ((b * C_OUT + c_out) * OH + oh) * num_ow_tiles + ow_tile`
393///
394/// For each `(c_in_local, kh, kw)` combination:
395///   - `dx`: `atomic_add(dx_ptr[b, c_in, ih, iw], dy * w)`
396///   - `dw`: `atomic_add(dw_ptr[c_out, c_in_local, kh, kw], sum(dy * x))`
397///
398/// Both `dx_ptr` and `dw_ptr` must be zero-initialised before launch.
399#[kernel]
400pub fn conv2d_backward<
401    T: Triton,
402    D: Num,
403    const KH: i32,
404    const KW: i32,
405    const STRIDE_H: i32,
406    const STRIDE_W: i32,
407    const PAD_H: i32,
408    const PAD_W: i32,
409    const G: i32,
410    const BLOCK_OW: i32,
411>(
412    dy_ptr: T::Pointer<D>,
413    x_ptr: T::Pointer<D>,
414    w_ptr: T::Pointer<D>,
415    dx_ptr: T::Pointer<D>,
416    dw_ptr: T::Pointer<D>,
417    _B: i32,
418    C_IN: i32,
419    C_OUT: i32,
420    H: i32,
421    W: i32,
422    OH: i32,
423    OW: i32,
424) where
425    T::I32Tensor: Tensor<i32, 1>,
426    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
427    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
428    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
429{
430    let pid = T::program_id(Axis::X);
431    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
432
433    let ow_tile = pid % num_ow_tiles;
434    let bco = pid / num_ow_tiles;
435    let oh = bco % OH;
436    let bc = bco / OH;
437    let c_out = bc % C_OUT;
438    let b = bc / C_OUT;
439
440    let ow_start = ow_tile * BLOCK_OW;
441    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
442    let ow_mask = ow_range.lt(OW);
443
444    // Load the upstream gradient tile for this output row.
445    let dy_offsets = ow_range + ((b * C_OUT + c_out) * OH * OW + oh * OW);
446    let dy_tile = T::load(
447        dy_ptr.add_offsets(dy_offsets),
448        Some(ow_mask),
449        Some(T::zeros::<D>(&[BLOCK_OW])),
450        &[],
451        None,
452        None,
453        None,
454        false,
455    );
456
457    let c_in_per_group = C_IN / G;
458    let g_idx = c_out / (C_OUT / G);
459    let c_in_start = g_idx * c_in_per_group;
460
461    let loop_bound = c_in_per_group * KH * KW;
462    for idx in 0..loop_bound {
463        let kw = idx % KW;
464        let kh_cin = idx / KW;
465        let kh = kh_cin % KH;
466        let c_in_local = kh_cin / KH;
467        let c_in = c_in_start + c_in_local;
468
469        // Load weight scalar and broadcast.
470        let w_idx = ((c_out * c_in_per_group + c_in_local) * KH + kh) * KW + kw;
471        let w_off = T::arange(0, 1) + w_idx;
472        let w_1 = T::load(
473            w_ptr.add_offsets(w_off),
474            None,
475            None,
476            &[],
477            None,
478            None,
479            None,
480            false,
481        );
482        let w_tile = T::broadcast_to(w_1, &[BLOCK_OW]);
483
484        let ih = oh * STRIDE_H + kh - PAD_H;
485        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
486
487        #[allow(clippy::erasing_op)]
488        let ih_t = ow_range * 0 + ih;
489        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
490        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
491        let in_mask = ow_mask & h_in_bounds & w_in_bounds;
492
493        let dx_offsets = iw_range + ((b * C_IN + c_in) * H * W + ih * W);
494
495        // dx: scatter dy * w to input positions.
496        let grad_tile = dy_tile * w_tile;
497        T::atomic_add(
498            dx_ptr.add_offsets(dx_offsets),
499            grad_tile,
500            Some(in_mask),
501            None,
502            None,
503        );
504
505        // dw: load x tile, compute partial = sum(dy * x), scatter to weight.
506        let x_tile = T::load(
507            x_ptr.add_offsets(dx_offsets),
508            Some(in_mask),
509            Some(T::zeros::<D>(&[BLOCK_OW])),
510            &[],
511            None,
512            None,
513            None,
514            false,
515        );
516        let partial = T::sum(dy_tile * x_tile, Some(0), false);
517        let partial_1 = T::expand_dims(partial, 0);
518        let dw_off = T::arange(0, 1) + w_idx;
519        T::atomic_add(dw_ptr.add_offsets(dw_off), partial_1, None, None, None);
520    }
521}
522
523impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for Conv2dForward<D> {
524    fn n_activation_inputs(&self) -> usize {
525        1
526    }
527
528    fn param_shapes(&self, input_shapes: &[&[usize]], output_shape: &[usize]) -> Vec<Vec<usize>> {
529        let c_in = input_shapes[0][1];
530        let c_out = output_shape[1];
531        vec![vec![
532            c_out,
533            c_in / self.g as usize,
534            self.kh as usize,
535            self.kw as usize,
536        ]]
537    }
538
539    fn param_names(&self) -> &'static [&'static str] {
540        &["weight"]
541    }
542
543    fn pack_args(
544        &self,
545        inputs: &[(teeny_core::model::RawPtr, &[usize])],
546        params: &[teeny_core::model::RawPtr],
547        output: teeny_core::model::RawPtr,
548        output_shape: &[usize],
549        _output_row_stride: i32,
550        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
551    ) {
552        let input_shape = inputs[0].1;
553        visitor.visit_ptr(inputs[0].0);
554        visitor.visit_ptr(params[0]);
555        visitor.visit_ptr(output);
556        visitor.visit_i32(input_shape[0] as i32);
557        visitor.visit_i32(input_shape[1] as i32);
558        visitor.visit_i32(output_shape[1] as i32);
559        visitor.visit_i32(input_shape[2] as i32);
560        visitor.visit_i32(input_shape[3] as i32);
561        visitor.visit_i32(output_shape[2] as i32);
562        visitor.visit_i32(output_shape[3] as i32);
563    }
564
565    fn block(&self) -> [u32; 3] {
566        [128, 1, 1]
567    }
568
569    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
570        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
571        [
572            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
573            1,
574            1,
575        ]
576    }
577
578    #[cfg(feature = "training")]
579    fn has_backward(&self) -> bool {
580        true
581    }
582
583    /// kernel args: dy, x, w, dx, dw, B, C_IN, C_OUT, H, W, OH, OW
584    #[cfg(feature = "training")]
585    fn pack_backward_args(
586        &self,
587        inputs: &[(teeny_core::model::RawPtr, &[usize])],
588        params: &[teeny_core::model::RawPtr],
589        _output: teeny_core::model::RawPtr,
590        output_shape: &[usize],
591        grad_output: teeny_core::model::RawPtr,
592        _grad_output_row_stride: i32,
593        grad_inputs: &[teeny_core::model::RawPtr],
594        grad_params: &[teeny_core::model::RawPtr],
595        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
596    ) {
597        let in_shape = inputs[0].1; // [B, C_IN, H, W]
598        visitor.visit_ptr(grad_output); // dy_ptr
599        visitor.visit_ptr(inputs[0].0); // x_ptr
600        visitor.visit_ptr(params[0]); // w_ptr
601        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
602        visitor.visit_ptr(grad_params[0]); // dw_ptr
603        visitor.visit_i32(in_shape[0] as i32); // B
604        visitor.visit_i32(in_shape[1] as i32); // C_IN
605        visitor.visit_i32(output_shape[1] as i32); // C_OUT
606        visitor.visit_i32(in_shape[2] as i32); // H
607        visitor.visit_i32(in_shape[3] as i32); // W
608        visitor.visit_i32(output_shape[2] as i32); // OH
609        visitor.visit_i32(output_shape[3] as i32); // OW
610    }
611
612    #[cfg(feature = "training")]
613    fn backward_block(&self) -> [u32; 3] {
614        [128, 1, 1]
615    }
616
617    /// Grid over forward-output positions (same formula as forward).
618    #[cfg(feature = "training")]
619    fn backward_grid(&self, _input_shapes: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
620        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
621        [
622            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
623            1,
624            1,
625        ]
626    }
627}
628
629pub struct Conv2dOp<'a, T: Num> {
630    pub forward: Conv2dForward<T>,
631    pub backward_dx: Conv2dBackwardDx<T>,
632    pub backward_dw: Conv2dBackwardDw<T>,
633    _marker: core::marker::PhantomData<&'a ()>,
634}
635
636/// 2-D convolution forward pass fused with a per-output-channel bias add.
637///
638/// Identical to [`conv2d_forward`] (same grid, same masked-load accumulation loop —
639/// see its doc comment) with one addition: `acc + bias[c_out]` before the store.
640/// `bias` broadcasts the same "load as a [1] tensor, then broadcast_to" pattern
641/// `conv2d_forward` already uses for weights, applied once after the loop instead
642/// of once per (c_in, kh, kw) tap.
643///
644/// For `Conv2d(has_bias=true)` with no downstream BatchNorm/SiLU to fuse into (the
645/// `conv2d_bn_silu` family), this replaces what would otherwise lower to two
646/// separate kernel launches — [`conv2d_forward`] then a standalone NCHW bias-add —
647/// with one. See spinorml-ia5.
648///
649/// Inference-only; no backward pass (training still uses the two-kernel path via
650/// `conv2d_forward` + a separate bias-add, whose backward is a plain per-channel sum
651/// over the output gradient — fusing that isn't this kernel's job).
652#[kernel]
653pub fn conv2d_bias_forward<
654    T: Triton,
655    D: Num,
656    const KH: i32,
657    const KW: i32,
658    const STRIDE_H: i32,
659    const STRIDE_W: i32,
660    const PAD_H: i32,
661    const PAD_W: i32,
662    const G: i32,
663    const BLOCK_OW: i32,
664>(
665    x_ptr: T::Pointer<D>,
666    w_ptr: T::Pointer<D>,
667    bias_ptr: T::Pointer<D>,
668    y_ptr: T::Pointer<D>,
669    _B: i32,
670    C_IN: i32,
671    C_OUT: i32,
672    H: i32,
673    W: i32,
674    OH: i32,
675    OW: i32,
676) where
677    T::I32Tensor: Tensor<i32, 1>,
678    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
679    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
680    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
681{
682    let pid = T::program_id(Axis::X);
683    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
684
685    let ow_tile = pid % num_ow_tiles;
686    let bco = pid / num_ow_tiles;
687    let oh = bco % OH;
688    let bc = bco / OH;
689    let c_out = bc % C_OUT;
690    let b = bc / C_OUT;
691
692    let ow_start = ow_tile * BLOCK_OW;
693    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
694    let ow_mask = ow_range.lt(OW);
695
696    let out_bc_base = (b * C_OUT + c_out) * OH * OW;
697
698    let c_in_per_group = C_IN / G;
699    let g_idx = c_out / (C_OUT / G);
700    let c_in_start = g_idx * c_in_per_group;
701
702    let mut acc = T::zeros::<D>(&[BLOCK_OW]);
703
704    let loop_bound = c_in_per_group * KH * KW;
705    for idx in 0..loop_bound {
706        let kw = idx % KW;
707        let kh_cin = idx / KW;
708        let kh = kh_cin % KH;
709        let c_in_local = kh_cin / KH;
710        let c_in = c_in_start + c_in_local;
711
712        let ih = oh * STRIDE_H + kh - PAD_H;
713        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
714
715        #[allow(clippy::erasing_op)]
716        let ih_t = ow_range * 0 + ih;
717        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
718        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
719        let load_mask = ow_mask & h_in_bounds & w_in_bounds;
720
721        let x_offsets = iw_range + ((b * C_IN + c_in) * H * W + ih * W);
722        let x_tile = T::load(
723            x_ptr.add_offsets(x_offsets),
724            Some(load_mask),
725            Some(T::zeros::<D>(&[BLOCK_OW])),
726            &[],
727            None,
728            None,
729            None,
730            false,
731        );
732
733        let w_idx = ((c_out * c_in_per_group + c_in_local) * KH + kh) * KW + kw;
734        let w_off = T::arange(0, 1) + w_idx;
735        let w_1 = T::load(
736            w_ptr.add_offsets(w_off),
737            None,
738            None,
739            &[],
740            None,
741            None,
742            None,
743            false,
744        );
745        let w_tile = T::broadcast_to(w_1, &[BLOCK_OW]);
746
747        acc = acc + x_tile * w_tile;
748    }
749
750    // ── Bias epilog ──────────────────────────────────────────────────────────
751    let bias_off = T::arange(0, 1) + c_out;
752    let bias_1 = T::load(
753        bias_ptr.add_offsets(bias_off),
754        None,
755        None,
756        &[],
757        None,
758        None,
759        None,
760        false,
761    );
762    let bias_tile = T::broadcast_to(bias_1, &[BLOCK_OW]);
763    acc = acc + bias_tile;
764
765    let out_offsets = ow_range + (out_bc_base + oh * OW);
766    T::store(
767        y_ptr.add_offsets(out_offsets),
768        acc,
769        Some(ow_mask),
770        &[],
771        None,
772        None,
773    );
774}
775
776impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for Conv2dBiasForward<D> {
777    fn n_activation_inputs(&self) -> usize {
778        1
779    }
780
781    fn param_shapes(&self, input_shapes: &[&[usize]], output_shape: &[usize]) -> Vec<Vec<usize>> {
782        let c_in = input_shapes[0][1];
783        let c_out = output_shape[1];
784        vec![
785            vec![c_out, c_in / self.g as usize, self.kh as usize, self.kw as usize],
786            vec![c_out],
787        ]
788    }
789
790    fn param_names(&self) -> &'static [&'static str] {
791        &["weight", "bias"]
792    }
793
794    fn pack_args(
795        &self,
796        inputs: &[(teeny_core::model::RawPtr, &[usize])],
797        params: &[teeny_core::model::RawPtr],
798        output: teeny_core::model::RawPtr,
799        output_shape: &[usize],
800        _output_row_stride: i32,
801        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
802    ) {
803        let input_shape = inputs[0].1;
804        visitor.visit_ptr(inputs[0].0); // x_ptr
805        visitor.visit_ptr(params[0]); // w_ptr
806        visitor.visit_ptr(params[1]); // bias_ptr
807        visitor.visit_ptr(output); // y_ptr
808        visitor.visit_i32(input_shape[0] as i32); // B
809        visitor.visit_i32(input_shape[1] as i32); // C_IN
810        visitor.visit_i32(output_shape[1] as i32); // C_OUT
811        visitor.visit_i32(input_shape[2] as i32); // H
812        visitor.visit_i32(input_shape[3] as i32); // W
813        visitor.visit_i32(output_shape[2] as i32); // OH
814        visitor.visit_i32(output_shape[3] as i32); // OW
815    }
816
817    fn block(&self) -> [u32; 3] {
818        [128, 1, 1]
819    }
820
821    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
822        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
823        [
824            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
825            1,
826            1,
827        ]
828    }
829}
teeny_kernels/nn/conv/conv2d.rs

teeny_kernels/nn/conv/
conv2d.rs
