teeny_kernels/nn/tensor/
channel_bias_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//! Channel bias add Triton kernels.
18//!
19//! Layout: input `x` and output `y` are NC-layout, where N = B*H*W and C =
20//! number of channels. Bias is a (C,) vector. Element `x[n, c]` lives at flat
21//! offset `n * C + c`.
22//!
23//! Parallelism: **one CTA per channel**. Each CTA iterates over all N spatial
24//! elements in `BLOCK_N`-wide tiles, adding the per-channel scalar bias.
25
26#![allow(non_snake_case)]
27
28use teeny_core::dtype::Float;
29use teeny_macros::kernel;
30use teeny_triton::triton::{
31    types::{AddOffsets, Comparison},
32    *,
33};
34
35// ─── Forward ─────────────────────────────────────────────────────────────────
36
37/// Adds a (C,) bias to a tensor in NC layout (N = B*H*W, C = channels).
38///
39/// Grid: `[C]` — one CTA per channel.
40#[kernel]
41pub fn channel_bias_add_forward<T: Triton, D: Float, const BLOCK_N: i32>(
42    x_ptr: T::Pointer<D>,
43    bias_ptr: T::Pointer<D>,
44    y_ptr: T::Pointer<D>,
45    N: i32,
46    C: i32,
47) where
48    T::I32Tensor: types::Tensor<i32, 1>,
49    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
50    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
51{
52    let c = T::program_id(Axis::X);
53    let c_idx = T::arange(0, 1) + c;
54
55    // Load bias[c] as shape [1], broadcast to [BLOCK_N].
56    let bias = T::broadcast_to(
57        T::load(
58            bias_ptr.add_offsets(c_idx),
59            None,
60            None,
61            &[],
62            None,
63            None,
64            None,
65            false,
66        ),
67        &[BLOCK_N],
68    );
69
70    let zeros = T::zeros::<D>(&[BLOCK_N]);
71    let mut n_start: i32 = 0;
72    while n_start < N {
73        let offsets_n = T::arange(0, BLOCK_N) + n_start;
74        let mask = offsets_n.lt(N);
75        let elem_offsets = offsets_n * C + c;
76
77        let x_tile = T::load(
78            x_ptr.add_offsets(elem_offsets),
79            Some(mask),
80            Some(zeros),
81            &[],
82            None,
83            None,
84            None,
85            false,
86        );
87        T::store(
88            y_ptr.add_offsets(elem_offsets),
89            x_tile + bias,
90            Some(mask),
91            &[],
92            None,
93            None,
94        );
95
96        n_start += BLOCK_N;
97    }
98}
99
100// ─── Backward ────────────────────────────────────────────────────────────────
101
102/// Backward pass for channel bias add.
103///
104/// dx = dy (identity), dbias[c] = sum over N of dy[n, c].
105///
106/// Grid: `[C]` — one CTA per channel.
107#[kernel]
108pub fn channel_bias_add_backward<T: Triton, D: Float, const BLOCK_N: i32>(
109    dy_ptr: T::Pointer<D>,
110    dx_ptr: T::Pointer<D>,
111    dbias_ptr: T::Pointer<D>,
112    N: i32,
113    C: i32,
114) where
115    T::I32Tensor: types::Tensor<i32, 1>,
116    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
117    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
118{
119    let c = T::program_id(Axis::X);
120    let c_idx = T::arange(0, 1) + c;
121
122    let zeros = T::zeros::<D>(&[BLOCK_N]);
123    let mut acc = T::zeros::<D>(&[1]);
124    let mut n_start: i32 = 0;
125
126    while n_start < N {
127        let offsets_n = T::arange(0, BLOCK_N) + n_start;
128        let mask = offsets_n.lt(N);
129        let elem_offsets = offsets_n * C + c;
130
131        let dy_tile = T::load(
132            dy_ptr.add_offsets(elem_offsets),
133            Some(mask),
134            Some(zeros),
135            &[],
136            None,
137            None,
138            None,
139            false,
140        );
141
142        // dx = dy (pass-through gradient)
143        T::store(
144            dx_ptr.add_offsets(elem_offsets),
145            dy_tile,
146            Some(mask),
147            &[],
148            None,
149            None,
150        );
151
152        // Accumulate bias gradient (keepdim so shape stays [1])
153        acc = acc + T::sum(dy_tile, None, true);
154
155        n_start += BLOCK_N;
156    }
157
158    // Write accumulated bias gradient with atomic add.
159    T::atomic_add(dbias_ptr.add_offsets(c_idx), acc, None, None, None);
160}
161
162// ─── NCHW Bias Add (for Conv2d with bias) ────────────────────────────────────
163
164/// Adds a (C,) bias to a tensor in NCHW layout.
165///
166/// Grid: `[C, B]` — one CTA per (channel, batch) pair; each CTA iterates over
167/// H*W spatial positions in `BLOCK_HW`-wide tiles.
168#[kernel]
169pub fn nchw_bias_add_forward<T: Triton, D: Float, const BLOCK_HW: i32>(
170    x_ptr: T::Pointer<D>,
171    bias_ptr: T::Pointer<D>,
172    y_ptr: T::Pointer<D>,
173    C: i32,
174    HW: i32,
175) where
176    T::I32Tensor: types::Tensor<i32, 1>,
177    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
178    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
179{
180    let c = T::program_id(Axis::X);
181    let b = T::program_id(Axis::Y);
182    let c_idx = T::arange(0, 1) + c;
183
184    let bias = T::broadcast_to(
185        T::load(
186            bias_ptr.add_offsets(c_idx),
187            None,
188            None,
189            &[],
190            None,
191            None,
192            None,
193            false,
194        ),
195        &[BLOCK_HW],
196    );
197
198    let zeros = T::zeros::<D>(&[BLOCK_HW]);
199    let batch_channel_offset: i32 = b * C * HW + c * HW;
200    let mut hw_start: i32 = 0;
201    while hw_start < HW {
202        let offsets = T::arange(0, BLOCK_HW) + hw_start;
203        let mask = offsets.lt(HW);
204        let elem_offsets = offsets + batch_channel_offset;
205        let x_tile = T::load(
206            x_ptr.add_offsets(elem_offsets),
207            Some(mask),
208            Some(zeros),
209            &[],
210            None,
211            None,
212            None,
213            false,
214        );
215        T::store(
216            y_ptr.add_offsets(elem_offsets),
217            x_tile + bias,
218            Some(mask),
219            &[],
220            None,
221            None,
222        );
223        hw_start += BLOCK_HW;
224    }
225}
226
227/// NCHW bias add backward: dx = dy, dbias[c] = sum over (B, H, W) of dy.
228///
229/// Grid: `[C, B]` — one CTA per (channel, batch) pair; single while loop over
230/// H*W to avoid nested loops (which ICE the teenyc compiler).
231#[kernel]
232pub fn nchw_bias_add_backward<T: Triton, D: Float, const BLOCK_HW: i32>(
233    dy_ptr: T::Pointer<D>,
234    dx_ptr: T::Pointer<D>,
235    dbias_ptr: T::Pointer<D>,
236    C: i32,
237    HW: i32,
238) where
239    T::I32Tensor: types::Tensor<i32, 1>,
240    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
241    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
242{
243    let c = T::program_id(Axis::X);
244    let b = T::program_id(Axis::Y);
245    let c_idx = T::arange(0, 1) + c;
246    let zeros = T::zeros::<D>(&[BLOCK_HW]);
247    let mut dbias_acc = T::zeros::<D>(&[1]);
248    let batch_channel_offset: i32 = b * C * HW + c * HW;
249
250    let mut hw_start: i32 = 0;
251    while hw_start < HW {
252        let offsets = T::arange(0, BLOCK_HW) + hw_start;
253        let mask = offsets.lt(HW);
254        let elem_offsets = offsets + batch_channel_offset;
255        let dy_tile = T::load(
256            dy_ptr.add_offsets(elem_offsets),
257            Some(mask),
258            Some(zeros),
259            &[],
260            None,
261            None,
262            None,
263            false,
264        );
265        T::store(
266            dx_ptr.add_offsets(elem_offsets),
267            dy_tile,
268            Some(mask),
269            &[],
270            None,
271            None,
272        );
273        dbias_acc = dbias_acc + T::sum(dy_tile, None, true);
274        hw_start += BLOCK_HW;
275    }
276    T::atomic_add(dbias_ptr.add_offsets(c_idx), dbias_acc, None, None, None);
277}
278
279/// RuntimeOp for adding a (C,) bias to an NCHW-layout tensor.
280///
281/// Used by the graph lowering when `Op::Conv2d { has_bias: true }` is encountered.
282/// Grid: `[C, B]` — one CTA per (channel, batch-item).
283pub struct NchwBiasAddRuntimeOp<D: Float + Send + Sync + 'static> {
284    fwd: NchwBiasAddForward<D>,
285    bwd: NchwBiasAddBackward<D>,
286    block_hw: i32,
287}
288
289impl<D: Float + Send + Sync + 'static> NchwBiasAddRuntimeOp<D> {
290    pub fn new(block_hw: i32) -> Self {
291        Self {
292            fwd: NchwBiasAddForward::<D>::new(block_hw),
293            bwd: NchwBiasAddBackward::<D>::new(block_hw),
294            block_hw,
295        }
296    }
297    pub fn forward_source(&self) -> &str {
298        &self.fwd.source
299    }
300    pub fn backward_source(&self) -> &str {
301        &self.bwd.source
302    }
303    pub fn kernel_name(&self) -> &str {
304        self.fwd.name
305    }
306}
307
308impl<D: Float + Send + Sync + 'static> teeny_core::model::RuntimeOp for NchwBiasAddRuntimeOp<D> {
309    fn n_activation_inputs(&self) -> usize {
310        1
311    }
312
313    fn param_shapes(&self, input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
314        let c = input_shapes[0][1];
315        vec![vec![c]]
316    }
317
318    fn param_names(&self) -> &'static [&'static str] {
319        &["bias"]
320    }
321
322    fn pack_args(
323        &self,
324        inputs: &[(teeny_core::model::RawPtr, &[usize])],
325        params: &[teeny_core::model::RawPtr],
326        output: teeny_core::model::RawPtr,
327        output_shape: &[usize],
328        _output_row_stride: i32,
329        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
330    ) {
331        let c = output_shape[1] as i32;
332        let hw = (output_shape[2] * output_shape[3]) as i32;
333        visitor.visit_ptr(inputs[0].0); // x_ptr
334        visitor.visit_ptr(params[0]); // bias_ptr
335        visitor.visit_ptr(output); // y_ptr
336        visitor.visit_i32(c);
337        visitor.visit_i32(hw);
338    }
339
340    fn block(&self) -> [u32; 3] {
341        [self.block_hw as u32, 1, 1]
342    }
343
344    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
345        [output_shape[1] as u32, output_shape[0] as u32, 1]
346    }
347
348    #[cfg(feature = "training")]
349    fn has_backward(&self) -> bool {
350        true
351    }
352
353    #[cfg(feature = "training")]
354    fn pack_backward_args(
355        &self,
356        inputs: &[(teeny_core::model::RawPtr, &[usize])],
357        _params: &[teeny_core::model::RawPtr],
358        _output: teeny_core::model::RawPtr,
359        _output_shape: &[usize],
360        grad_output: teeny_core::model::RawPtr,
361        _grad_output_row_stride: i32,
362        grad_inputs: &[teeny_core::model::RawPtr],
363        grad_params: &[teeny_core::model::RawPtr],
364        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
365    ) {
366        let in_shape = inputs[0].1; // [B, C, H, W]
367        let c = in_shape[1] as i32;
368        let hw = (in_shape[2] * in_shape[3]) as i32;
369        visitor.visit_ptr(grad_output); // dy_ptr
370        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
371        visitor.visit_ptr(grad_params[0]); // dbias_ptr
372        // B is encoded in grid.y — not passed as a kernel arg.
373        visitor.visit_i32(c);
374        visitor.visit_i32(hw);
375    }
376
377    #[cfg(feature = "training")]
378    fn backward_block(&self) -> [u32; 3] {
379        [self.block_hw as u32, 1, 1]
380    }
381
382    #[cfg(feature = "training")]
383    fn backward_grid(&self, input_shapes: &[&[usize]], _output_shape: &[usize]) -> [u32; 3] {
384        // Grid [C, B] — one CTA per (channel, batch) pair, mirroring the kernel.
385        [input_shapes[0][1] as u32, input_shapes[0][0] as u32, 1]
386    }
387}
388
389// ─── RuntimeOp ───────────────────────────────────────────────────────────────
390
391/// Combined forward + backward RuntimeOp for channel bias add.
392///
393/// Forward kernel arguments: `x_ptr`, `bias_ptr`, `y_ptr`, `N_SPATIAL`, `C`.
394/// Backward kernel arguments: `dy_ptr`, `dx_ptr`, `dbias_ptr`, `N_SPATIAL`, `C`.
395///
396/// Grid: `[C, 1, 1]` for both forward and backward.
397pub struct ChannelBiasAddRuntimeOp<D: Float + Send + Sync + 'static> {
398    fwd: ChannelBiasAddForward<D>,
399    bwd: ChannelBiasAddBackward<D>,
400    /// Output channel count, fixed at construction time.
401    c_out: usize,
402}
403
404impl<D: Float + Send + Sync + 'static> ChannelBiasAddRuntimeOp<D> {
405    pub fn new(block_n: i32, c_out: usize) -> Self {
406        Self {
407            fwd: ChannelBiasAddForward::<D>::new(block_n),
408            bwd: ChannelBiasAddBackward::<D>::new(block_n),
409            c_out,
410        }
411    }
412
413    pub fn forward_source(&self) -> &str {
414        &self.fwd.source
415    }
416    pub fn backward_source(&self) -> &str {
417        &self.bwd.source
418    }
419    pub fn kernel_name(&self) -> &str {
420        self.fwd.name
421    }
422}
423
424impl<D: Float + Send + Sync + 'static> teeny_core::model::RuntimeOp for ChannelBiasAddRuntimeOp<D> {
425    fn n_activation_inputs(&self) -> usize {
426        1
427    }
428
429    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
430        // Bias shape is (C_out,)
431        vec![vec![self.c_out]]
432    }
433
434    fn param_names(&self) -> &'static [&'static str] {
435        &["bias"]
436    }
437
438    fn pack_args(
439        &self,
440        inputs: &[(teeny_core::model::RawPtr, &[usize])],
441        params: &[teeny_core::model::RawPtr],
442        output: teeny_core::model::RawPtr,
443        output_shape: &[usize],
444        _output_row_stride: i32,
445        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
446    ) {
447        // inputs[0].1 = [B, C, H, W]
448        let input_shape = inputs[0].1;
449        let b = input_shape[0];
450        let c = output_shape[1];
451        let h = input_shape[2];
452        let w = input_shape[3];
453        let n_spatial = (b * h * w) as i32;
454
455        visitor.visit_ptr(inputs[0].0); // x_ptr
456        visitor.visit_ptr(params[0]); // bias_ptr
457        visitor.visit_ptr(output); // y_ptr
458        visitor.visit_i32(n_spatial); // N
459        visitor.visit_i32(c as i32); // C
460    }
461
462    fn block(&self) -> [u32; 3] {
463        [128, 1, 1]
464    }
465
466    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
467        // One CTA per output channel.
468        [output_shape[1] as u32, 1, 1]
469    }
470
471    #[cfg(feature = "training")]
472    fn has_backward(&self) -> bool {
473        true
474    }
475
476    #[cfg(feature = "training")]
477    fn pack_backward_args(
478        &self,
479        inputs: &[(teeny_core::model::RawPtr, &[usize])],
480        _params: &[teeny_core::model::RawPtr],
481        _output: teeny_core::model::RawPtr,
482        output_shape: &[usize],
483        grad_output: teeny_core::model::RawPtr,
484        _grad_output_row_stride: i32,
485        grad_inputs: &[teeny_core::model::RawPtr],
486        grad_params: &[teeny_core::model::RawPtr],
487        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
488    ) {
489        let input_shape = inputs[0].1;
490        let b = input_shape[0];
491        let c = output_shape[1];
492        let h = input_shape[2];
493        let w = input_shape[3];
494        let n_spatial = (b * h * w) as i32;
495
496        visitor.visit_ptr(grad_output); // dy_ptr
497        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
498        visitor.visit_ptr(grad_params[0]); // dbias_ptr
499        visitor.visit_i32(n_spatial); // N
500        visitor.visit_i32(c as i32); // C
501    }
502
503    #[cfg(feature = "training")]
504    fn backward_grid(&self, _input_shapes: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
505        [output_shape[1] as u32, 1, 1]
506    }
507}
teeny_kernels/nn/tensor/channel_bias_add.rs

teeny_kernels/nn/tensor/
channel_bias_add.rs
