teeny_kernels/nn/conv/
conv3d.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/// 3-D convolution forward pass.
29///
30/// Grid: one CTA per (b, c_out, od, oh, ow-tile):
31///   `pid = (((b * C_OUT + c_out) * OD + od) * OH + oh) * num_ow_tiles + ow_tile`
32///
33/// Each CTA computes a BLOCK_OW-wide strip of output width positions by
34/// iterating over all `C_IN * KD * KH * KW` combinations.
35///
36/// Zero-padding of `PAD_D`/`PAD_H`/`PAD_W` elements is applied on each side.
37/// `OD = (D + 2*PAD_D - KD) / STRIDE_D + 1`, etc.
38#[kernel]
39pub fn conv3d_forward<
40    T: Triton,
41    D: Num,
42    const KD: i32,
43    const KH: i32,
44    const KW: i32,
45    const STRIDE_D: i32,
46    const STRIDE_H: i32,
47    const STRIDE_W: i32,
48    const PAD_D: i32,
49    const PAD_H: i32,
50    const PAD_W: i32,
51    const BLOCK_OW: i32,
52>(
53    x_ptr: T::Pointer<D>,
54    w_ptr: T::Pointer<D>,
55    y_ptr: T::Pointer<D>,
56    _B: i32,
57    C_IN: i32,
58    C_OUT: i32,
59    Dv: i32,
60    H: i32,
61    W: i32,
62    OD: i32,
63    OH: i32,
64    OW: i32,
65) where
66    T::I32Tensor: Tensor<i32, 1>,
67    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
68    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
69    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
70{
71    let pid = T::program_id(Axis::X);
72    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
73
74    // Decode flat pid → (b, c_out, od, oh, ow_tile).
75    let ow_tile = pid % num_ow_tiles;
76    let rest = pid / num_ow_tiles;
77    let oh = rest % OH;
78    let rest2 = rest / OH;
79    let od = rest2 % OD;
80    let bco = rest2 / OD;
81    let c_out = bco % C_OUT;
82    let b = bco / C_OUT;
83
84    let ow_start = ow_tile * BLOCK_OW;
85    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
86    let ow_mask = ow_range.lt(OW);
87
88    let out_base = ((b * C_OUT + c_out) * OD * OH * OW) + od * OH * OW + oh * OW;
89
90    let mut acc = T::zeros::<D>(&[BLOCK_OW]);
91
92    let loop_bound = C_IN * KD * KH * KW;
93    for idx in 0..loop_bound {
94        let kw = idx % KW;
95        let tmp = idx / KW;
96        let kh = tmp % KH;
97        let tmp2 = tmp / KH;
98        let kd = tmp2 % KD;
99        let c_in = tmp2 / KD;
100
101        // Compute padded input coordinates; OOB depth/height rows contribute zero via mask.
102        let id = od * STRIDE_D + kd - PAD_D;
103        let ih = oh * STRIDE_H + kh - PAD_H;
104        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
105
106        // `ow_range * 0` is the only way to splat scalar id/ih into an I32Tensor.
107        // Scalar `if`/`continue` inside a loop triggers a compiler phi-node bug.
108        #[allow(clippy::erasing_op)]
109        let id_t = ow_range * 0 + id;
110        #[allow(clippy::erasing_op)]
111        let ih_t = ow_range * 0 + ih;
112        let d_in_bounds = id_t.ge(0) & id_t.lt(Dv);
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 & d_in_bounds & h_in_bounds & w_in_bounds;
116
117        let x_offsets = iw_range + ((b * C_IN + c_in) * Dv * H * W + id * 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        let w_idx = (((c_out * C_IN + c_in) * KD + kd) * KH + kh) * KW + kw;
130        let w_off = T::arange(0, 1) + w_idx;
131        let w_1 = T::load(
132            w_ptr.add_offsets(w_off),
133            None,
134            None,
135            &[],
136            None,
137            None,
138            None,
139            false,
140        );
141        let w_tile = T::broadcast_to(w_1, &[BLOCK_OW]);
142
143        acc = acc + x_tile * w_tile;
144    }
145
146    let out_offsets = ow_range + out_base;
147    T::store(
148        y_ptr.add_offsets(out_offsets),
149        acc,
150        Some(ow_mask),
151        &[],
152        None,
153        None,
154    );
155}
156
157/// 3-D convolution backward pass — gradient with respect to input (`dx`).
158///
159/// Grid: `pid = (((b * C_OUT + c_out) * OD + od) * OH + oh) * num_ow_tiles + ow_tile`
160///
161/// Scatters gradient via `atomic_add` to handle overlapping receptive fields.
162/// Padding positions are skipped.
163#[kernel]
164pub fn conv3d_backward_dx<
165    T: Triton,
166    D: Num,
167    const KD: i32,
168    const KH: i32,
169    const KW: i32,
170    const STRIDE_D: i32,
171    const STRIDE_H: i32,
172    const STRIDE_W: i32,
173    const PAD_D: i32,
174    const PAD_H: i32,
175    const PAD_W: i32,
176    const BLOCK_OW: i32,
177>(
178    dy_ptr: T::Pointer<D>,
179    w_ptr: T::Pointer<D>,
180    dx_ptr: T::Pointer<D>,
181    _B: i32,
182    C_IN: i32,
183    C_OUT: i32,
184    Dv: i32,
185    H: i32,
186    W: i32,
187    OD: i32,
188    OH: i32,
189    OW: i32,
190) where
191    T::I32Tensor: Tensor<i32, 1>,
192    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
193    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
194    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
195{
196    let pid = T::program_id(Axis::X);
197    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
198
199    let ow_tile = pid % num_ow_tiles;
200    let rest = pid / num_ow_tiles;
201    let oh = rest % OH;
202    let rest2 = rest / OH;
203    let od = rest2 % OD;
204    let bco = rest2 / OD;
205    let c_out = bco % C_OUT;
206    let b = bco / C_OUT;
207
208    let ow_start = ow_tile * BLOCK_OW;
209    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
210    let ow_mask = ow_range.lt(OW);
211
212    let dy_offsets = ow_range + ((b * C_OUT + c_out) * OD * OH * OW + od * OH * OW + oh * OW);
213    let dy_tile = T::load(
214        dy_ptr.add_offsets(dy_offsets),
215        Some(ow_mask),
216        Some(T::zeros::<D>(&[BLOCK_OW])),
217        &[],
218        None,
219        None,
220        None,
221        false,
222    );
223
224    let loop_bound = C_IN * KD * KH * KW;
225    for idx in 0..loop_bound {
226        let kw = idx % KW;
227        let tmp = idx / KW;
228        let kh = tmp % KH;
229        let tmp2 = tmp / KH;
230        let kd = tmp2 % KD;
231        let c_in = tmp2 / KD;
232
233        let w_idx = (((c_out * C_IN + c_in) * KD + kd) * KH + kh) * KW + kw;
234        let w_off = T::arange(0, 1) + w_idx;
235        let w_1 = T::load(
236            w_ptr.add_offsets(w_off),
237            None,
238            None,
239            &[],
240            None,
241            None,
242            None,
243            false,
244        );
245        let w_tile = T::broadcast_to(w_1, &[BLOCK_OW]);
246
247        let grad_tile = dy_tile * w_tile;
248
249        let id = od * STRIDE_D + kd - PAD_D;
250        let ih = oh * STRIDE_H + kh - PAD_H;
251        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
252
253        #[allow(clippy::erasing_op)]
254        let id_t = ow_range * 0 + id;
255        #[allow(clippy::erasing_op)]
256        let ih_t = ow_range * 0 + ih;
257        let d_in_bounds = id_t.ge(0) & id_t.lt(Dv);
258        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
259        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
260
261        let dx_offsets = iw_range + ((b * C_IN + c_in) * Dv * H * W + id * H * W + ih * W);
262        T::atomic_add(
263            dx_ptr.add_offsets(dx_offsets),
264            grad_tile,
265            Some(ow_mask & d_in_bounds & h_in_bounds & w_in_bounds),
266            None,
267            None,
268        );
269    }
270}
271
272/// 3-D convolution backward pass — gradient with respect to weights (`dw`).
273///
274/// Grid: `pid = (((b * C_OUT + c_out) * OD + od) * OH + oh) * num_ow_tiles + ow_tile`
275///
276/// `dw` must be zero-initialised before launch.
277#[kernel]
278pub fn conv3d_backward_dw<
279    T: Triton,
280    D: Num,
281    const KD: i32,
282    const KH: i32,
283    const KW: i32,
284    const STRIDE_D: i32,
285    const STRIDE_H: i32,
286    const STRIDE_W: i32,
287    const PAD_D: i32,
288    const PAD_H: i32,
289    const PAD_W: i32,
290    const BLOCK_OW: i32,
291>(
292    dy_ptr: T::Pointer<D>,
293    x_ptr: T::Pointer<D>,
294    dw_ptr: T::Pointer<D>,
295    _B: i32,
296    C_IN: i32,
297    C_OUT: i32,
298    Dv: i32,
299    H: i32,
300    W: i32,
301    OD: i32,
302    OH: i32,
303    OW: i32,
304) where
305    T::I32Tensor: Tensor<i32, 1>,
306    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
307    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
308    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
309{
310    let pid = T::program_id(Axis::X);
311    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
312
313    let ow_tile = pid % num_ow_tiles;
314    let rest = pid / num_ow_tiles;
315    let oh = rest % OH;
316    let rest2 = rest / OH;
317    let od = rest2 % OD;
318    let bco = rest2 / OD;
319    let c_out = bco % C_OUT;
320    let b = bco / C_OUT;
321
322    let ow_start = ow_tile * BLOCK_OW;
323    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
324    let ow_mask = ow_range.lt(OW);
325
326    let dy_offsets = ow_range + ((b * C_OUT + c_out) * OD * OH * OW + od * OH * OW + oh * OW);
327    let dy_tile = T::load(
328        dy_ptr.add_offsets(dy_offsets),
329        Some(ow_mask),
330        Some(T::zeros::<D>(&[BLOCK_OW])),
331        &[],
332        None,
333        None,
334        None,
335        false,
336    );
337
338    let id_base = od * STRIDE_D;
339    let ih_base = oh * STRIDE_H;
340
341    let loop_bound = C_IN * KD * KH * KW;
342    for idx in 0..loop_bound {
343        let kw = idx % KW;
344        let tmp = idx / KW;
345        let kh = tmp % KH;
346        let tmp2 = tmp / KH;
347        let kd = tmp2 % KD;
348        let c_in = tmp2 / KD;
349
350        let id = id_base + kd - PAD_D;
351        let ih = ih_base + kh - PAD_H;
352        let iw_range = ow_range * STRIDE_W + kw - PAD_W;
353
354        #[allow(clippy::erasing_op)]
355        let id_t = ow_range * 0 + id;
356        #[allow(clippy::erasing_op)]
357        let ih_t = ow_range * 0 + ih;
358        let d_in_bounds = id_t.ge(0) & id_t.lt(Dv);
359        let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
360        let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
361        let load_mask = ow_mask & d_in_bounds & h_in_bounds & w_in_bounds;
362
363        let x_offsets = iw_range + ((b * C_IN + c_in) * Dv * H * W + id * H * W + ih * W);
364        let x_tile = T::load(
365            x_ptr.add_offsets(x_offsets),
366            Some(load_mask),
367            Some(T::zeros::<D>(&[BLOCK_OW])),
368            &[],
369            None,
370            None,
371            None,
372            false,
373        );
374
375        let partial = T::sum(dy_tile * x_tile, Some(0), false);
376        let partial_1 = T::expand_dims(partial, 0);
377
378        let w_idx = (((c_out * C_IN + c_in) * KD + kd) * KH + kh) * KW + kw;
379        let dw_off = T::arange(0, 1) + w_idx;
380        T::atomic_add(dw_ptr.add_offsets(dw_off), partial_1, None, None, None);
381    }
382}
383
384pub struct Conv3dOp<'a, T: Num> {
385    pub forward: Conv3dForward<T>,
386    pub backward_dx: Conv3dBackwardDx<T>,
387    pub backward_dw: Conv3dBackwardDw<T>,
388    _marker: core::marker::PhantomData<&'a ()>,
389}
teeny_kernels/nn/conv/conv3d.rs

teeny_kernels/nn/conv/
conv3d.rs
