vision_rs/models/yolo/kernels/loss/
cls.rs

1//! BCE classification loss Triton kernel for YOLO detection training.
2//!
3//! Computes numerically-stable binary cross-entropy loss between predicted
4//! class logits and soft target labels, summed over all C classes per anchor.
5//!
6//! Stable BCE formula (avoids log(0) and exp overflow):
7//!   loss = relu(x) − x·t + log(1 + exp(−|x|))
8//!        = max(x, 0) − x·t + log(1 + exp(−|x|))
9//!
10//! Layout:
11//!   - `pred`:   `[C, N]` — predicted class logits per anchor
12//!   - `target`: `[C, N]` — soft class labels ∈ [0, 1] per anchor
13//!   - `loss`:   `[N]`    — per-anchor BCE loss (summed over C classes)
14//!
15//! Parallelism: one CTA per BLOCK_N-wide anchor tile; classes iterated
16//! sequentially inside each CTA.
17//! Grid: `cdiv(N, BLOCK_N)` flat CTAs.
18
19#![allow(non_snake_case)]
20
21use teeny_core::dtype::Float;
22use teeny_macros::kernel;
23use teeny_triton::triton::{
24    types::{AddOffsets, Comparison},
25    *,
26};
27
28/// BCE classification loss forward: class logits + soft targets → per-anchor loss.
29///
30/// Grid: `cdiv(N, BLOCK_N)` — one CTA per anchor tile.
31#[allow(clippy::erasing_op, clippy::identity_op)]
32#[kernel]
33pub fn yolo_bce_cls_loss_forward<T: Triton, D: Float, const BLOCK_N: i32>(
34    pred_ptr: T::Pointer<D>,
35    target_ptr: T::Pointer<D>,
36    loss_ptr: T::Pointer<D>,
37    N: i32,
38    C: i32,
39) where
40    T::I32Tensor: types::Tensor<i32, 1>,
41    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
42    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
43{
44    let n_start = T::program_id(Axis::X) * BLOCK_N;
45    let n_offs  = T::arange(0, BLOCK_N) + n_start;
46    let mask    = n_offs.lt(N);
47    let zeros   = T::zeros::<D>(&[BLOCK_N]);
48    let ones    = T::full(&[BLOCK_N], D::from_f64(1.0));
49
50    // Accumulate BCE over all C classes for this anchor tile.
51    let mut acc = zeros;
52    let mut c: i32 = 0;
53    while c < C {
54        // Layout [C, N]: class c occupies row c, at base offset c*N.
55        let base = c * N;
56        let x = T::load(pred_ptr.add_offsets(n_offs + base),   Some(mask), Some(zeros), &[], None, None, None, false);
57        let t = T::load(target_ptr.add_offsets(n_offs + base), Some(mask), Some(zeros), &[], None, None, None, false);
58
59        // Numerically-stable BCE: max(x, 0) − x·t + log(1 + exp(−|x|))
60        let relu_x    = T::maximum(x, zeros);
61        let log1p_exp = T::log(ones + T::exp(zeros - T::abs(x)));
62        let bce       = relu_x - x * t + log1p_exp;
63
64        acc = acc + bce;
65        c += 1;
66    }
67
68    T::store(loss_ptr.add_offsets(n_offs), acc, Some(mask), &[], None, None);
69}
70
71/// BCE classification loss backward: upstream gradient + forward inputs → ∂L/∂pred.
72///
73/// Per element: `d_pred[c,n] = dy[n] · (sigmoid(pred[c,n]) − target[c,n])`
74///
75/// Uses the numerically stable sigmoid: `σ(x) = 1 / (1 + exp(−x))`.
76/// No saved activations are needed because the gradient depends only on the
77/// forward inputs (pred logits and targets).
78///
79/// Grid: `cdiv(N, BLOCK_N)` — one CTA per anchor tile.
80#[allow(clippy::erasing_op, clippy::identity_op)]
81#[kernel]
82pub fn yolo_bce_cls_loss_backward<T: Triton, D: Float, const BLOCK_N: i32>(
83    dy_ptr:     T::Pointer<D>,
84    pred_ptr:   T::Pointer<D>,
85    target_ptr: T::Pointer<D>,
86    d_pred_ptr: T::Pointer<D>,
87    N: i32,
88    C: i32,
89) where
90    T::I32Tensor: types::Tensor<i32, 1>,
91    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
92    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
93{
94    let n_start = T::program_id(Axis::X) * BLOCK_N;
95    let n_offs  = T::arange(0, BLOCK_N) + n_start;
96    let mask    = n_offs.lt(N);
97    let zeros   = T::zeros::<D>(&[BLOCK_N]);
98    let ones    = T::full(&[BLOCK_N], D::from_f64(1.0));
99    let neg_one = T::full(&[BLOCK_N], D::from_f64(-1.0));
100
101    let dy = T::load(dy_ptr.add_offsets(n_offs), Some(mask), Some(zeros), &[], None, None, None, false);
102
103    let mut c: i32 = 0;
104    while c < C {
105        let base = c * N;
106        let x = T::load(pred_ptr.add_offsets(n_offs + base),   Some(mask), Some(zeros), &[], None, None, None, false);
107        let t = T::load(target_ptr.add_offsets(n_offs + base), Some(mask), Some(zeros), &[], None, None, None, false);
108
109        // σ(x) = 1 / (1 + exp(−x))
110        let sig = ones / (ones + T::exp(neg_one * x));
111        let grad = dy * (sig - t);
112
113        T::store(d_pred_ptr.add_offsets(n_offs + base), grad, Some(mask), &[], None, None);
114        c += 1;
115    }
116}
vision_rs/models/yolo/kernels/loss/cls.rs

vision_rs/models/yolo/kernels/loss/
cls.rs
