teeny_kernels/nn/loss/
nll.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 teeny_macros::kernel;
20use teeny_triton::triton::{
21    types::{AddOffsets, Comparison},
22    *,
23};
24
25// ── NLLLoss ───────────────────────────────────────────────────────────────────
26
27/// NLL loss forward: `out[n] = -log_prob[n, target[n]]`.
28///
29/// Grid: `[n_rows, 1, 1]` — one CTA per batch element.
30///
31/// The flat index `pid * n_cols + target` is computed entirely in
32/// `T::Tensor<i32>` space (from T::full + T::load arithmetic), avoiding any
33/// conversion to `T::I32Tensor`.  A separate `AddOffsets` bound on
34/// `T::Pointer<_>` for `T::Tensor<i32>` covers the indexed load and store.
35#[kernel]
36pub fn nll_loss_forward<T: Triton>(
37    log_probs_ptr: T::Pointer<f32>,
38    targets_ptr: T::Pointer<i32>,
39    out_ptr: T::Pointer<f32>,
40    _n_rows: i32,
41    n_cols: i32,
42) where
43    T::I32Tensor: types::Tensor<i32, 1>,
44    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
45    // For loading targets (offset = T::I32Tensor from arange)
46    T::Pointer<i32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<i32>>>,
47    // For storing output (offset = T::I32Tensor from arange)
48    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
49    // For indexed load of log_prob (offset = T::Tensor<i32> from full + load arithmetic)
50    T::Tensor<i32>: types::Tensor<i32, 1>,
51    T::Pointer<f32>: AddOffsets<i32, 1, T::Tensor<i32>, Output = T::Tensor<T::Pointer<f32>>>,
52{
53    let pid = T::program_id(Axis::X);
54
55    // Load target class for this row (unmasked — pid is always a valid row index)
56    let tgt_off: T::I32Tensor = T::arange(0, 1) + pid;
57    let tgt: T::Tensor<i32> = T::load(
58        targets_ptr.add_offsets(tgt_off),
59        None,
60        None,
61        &[],
62        None,
63        None,
64        None,
65        false,
66    );
67
68    // flat index: pid * n_cols + target — purely in T::Tensor<i32> space
69    let base: T::Tensor<i32> = T::full::<i32>(&[1], pid * n_cols);
70    let flat_off: T::Tensor<i32> = base + tgt;
71
72    // Load log_prob at target class (unmasked — index is always valid)
73    let lp: T::Tensor<f32> = T::load(
74        log_probs_ptr.add_offsets(flat_off),
75        None,
76        None,
77        &[],
78        None,
79        None,
80        None,
81        false,
82    );
83
84    let loss = T::full(&[1], -1.0_f32) * lp;
85
86    let out_off: T::I32Tensor = T::arange(0, 1) + pid;
87    T::store(out_ptr.add_offsets(out_off), loss, None, &[], None, None);
88}
89
90/// NLL loss backward: `dx[n, target[n]] = -dy[n]`, zero elsewhere.
91///
92/// Grid: `[n_rows, 1, 1]` — one CTA per batch element.
93/// The dx buffer must be zero-initialised before launch.
94#[kernel]
95pub fn nll_loss_backward<T: Triton>(
96    dy_ptr: T::Pointer<f32>,
97    targets_ptr: T::Pointer<i32>,
98    dx_ptr: T::Pointer<f32>,
99    _n_rows: i32,
100    n_cols: i32,
101) where
102    T::I32Tensor: types::Tensor<i32, 1>,
103    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
104    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
105    T::Pointer<i32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<i32>>>,
106    T::Tensor<i32>: types::Tensor<i32, 1>,
107    T::Pointer<f32>: AddOffsets<i32, 1, T::Tensor<i32>, Output = T::Tensor<T::Pointer<f32>>>,
108{
109    let pid = T::program_id(Axis::X);
110
111    let dy_off: T::I32Tensor = T::arange(0, 1) + pid;
112    let dy: T::Tensor<f32> = T::load(
113        dy_ptr.add_offsets(dy_off),
114        None,
115        None,
116        &[],
117        None,
118        None,
119        None,
120        false,
121    );
122
123    let tgt_off: T::I32Tensor = T::arange(0, 1) + pid;
124    let tgt: T::Tensor<i32> = T::load(
125        targets_ptr.add_offsets(tgt_off),
126        None,
127        None,
128        &[],
129        None,
130        None,
131        None,
132        false,
133    );
134
135    // flat index: pid * n_cols + target
136    let base: T::Tensor<i32> = T::full::<i32>(&[1], pid * n_cols);
137    let flat_off: T::Tensor<i32> = base + tgt;
138
139    let neg_dy = T::full(&[1], -1.0_f32) * dy;
140    T::store(dx_ptr.add_offsets(flat_off), neg_dy, None, &[], None, None);
141}
142
143// ── CrossEntropyLoss ──────────────────────────────────────────────────────────
144
145/// Cross-entropy loss forward: `out[n] = log(sum_c exp(x[n,c])) - x[n, target[n]]`.
146///
147/// Numerically stable: subtracts row-max before exp (log-sum-exp trick).
148///
149/// Grid: `[n_rows, 1, 1]` — one CTA per row.
150/// `BLOCK_SIZE` must equal `next_power_of_two(n_cols)`.
151#[kernel]
152pub fn cross_entropy_loss_forward<T: Triton, const BLOCK_SIZE: i32>(
153    input_ptr: T::Pointer<f32>,
154    targets_ptr: T::Pointer<i32>,
155    out_ptr: T::Pointer<f32>,
156    _n_rows: i32,
157    n_cols: i32,
158) where
159    T::I32Tensor: types::Tensor<i32, 1>,
160    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
161    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
162    T::Pointer<i32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<i32>>>,
163    T::Tensor<i32>: types::Tensor<i32, 1>,
164    T::Pointer<f32>: AddOffsets<i32, 1, T::Tensor<i32>, Output = T::Tensor<T::Pointer<f32>>>,
165{
166    let pid = T::program_id(Axis::X);
167    let row_base = pid * n_cols;
168    let col_offs: T::I32Tensor = T::arange(0, BLOCK_SIZE);
169    let row_offs: T::I32Tensor = col_offs + row_base;
170    let in_row = col_offs.lt(n_cols);
171    // Large finite negative used as "neg-inf" mask: exp underflows to 0
172    let neg_inf = T::full(&[BLOCK_SIZE], -3.4028235e38_f32);
173
174    let row = T::load(
175        input_ptr.add_offsets(row_offs),
176        Some(in_row),
177        Some(neg_inf),
178        &[],
179        None,
180        None,
181        None,
182        false,
183    );
184
185    // log-sum-exp: subtract row max for numerical stability
186    let row_max = T::max(row, Some(0), true); // shape [1]
187    let row_shifted = row - row_max; // broadcast
188    let exp_row = T::exp(row_shifted);
189    let sum_exp = T::sum(exp_row, Some(0), true); // shape [1]
190    let log_sum_exp = T::log(sum_exp) + row_max; // shape [1]
191
192    // Load target index (unmasked — always in-bounds)
193    let tgt_off: T::I32Tensor = T::arange(0, 1) + pid;
194    let tgt: T::Tensor<i32> = T::load(
195        targets_ptr.add_offsets(tgt_off),
196        None,
197        None,
198        &[],
199        None,
200        None,
201        None,
202        false,
203    );
204
205    // flat index for x[target]: pid * n_cols + target
206    let base: T::Tensor<i32> = T::full::<i32>(&[1], row_base);
207    let flat_off: T::Tensor<i32> = base + tgt;
208    let x_target: T::Tensor<f32> = T::load(
209        input_ptr.add_offsets(flat_off),
210        None,
211        None,
212        &[],
213        None,
214        None,
215        None,
216        false,
217    );
218
219    // CE = log_sum_exp - x[target]
220    let loss = log_sum_exp - x_target;
221
222    let out_off: T::I32Tensor = T::arange(0, 1) + pid;
223    T::store(out_ptr.add_offsets(out_off), loss, None, &[], None, None);
224}
225
226/// Cross-entropy loss backward.
227///
228/// `dx[n, c] = dy[n] * (softmax(x[n])[c] - indicator(c == target[n]))`
229///
230/// Grid: `[n_rows, 1, 1]` — one CTA per row.
231/// `BLOCK_SIZE` must equal `next_power_of_two(n_cols)`.
232#[kernel]
233pub fn cross_entropy_loss_backward<T: Triton, const BLOCK_SIZE: i32>(
234    dy_ptr: T::Pointer<f32>,
235    input_ptr: T::Pointer<f32>,
236    targets_ptr: T::Pointer<i32>,
237    dx_ptr: T::Pointer<f32>,
238    _n_rows: i32,
239    n_cols: i32,
240) where
241    T::I32Tensor: types::Tensor<i32, 1>,
242    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
243    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
244    T::Pointer<i32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<i32>>>,
245    T::Tensor<i32>: types::Tensor<i32, 1>,
246    T::Pointer<f32>: AddOffsets<i32, 1, T::Tensor<i32>, Output = T::Tensor<T::Pointer<f32>>>,
247{
248    let pid = T::program_id(Axis::X);
249    let row_base = pid * n_cols;
250    let col_offs: T::I32Tensor = T::arange(0, BLOCK_SIZE);
251    let row_offs: T::I32Tensor = col_offs + row_base;
252    let in_row = col_offs.lt(n_cols);
253
254    // Load upstream gradient (unmasked — always in-bounds)
255    let dy_off: T::I32Tensor = T::arange(0, 1) + pid;
256    let dy: T::Tensor<f32> = T::load(
257        dy_ptr.add_offsets(dy_off),
258        None,
259        None,
260        &[],
261        None,
262        None,
263        None,
264        false,
265    );
266
267    // Load the row and compute softmax
268    let neg_inf = T::full(&[BLOCK_SIZE], -3.4028235e38_f32);
269    let row = T::load(
270        input_ptr.add_offsets(row_offs),
271        Some(in_row),
272        Some(neg_inf),
273        &[],
274        None,
275        None,
276        None,
277        false,
278    );
279    let sm = T::softmax(row, None, false, false);
280
281    // Load target index (unmasked — always in-bounds)
282    let tgt_off: T::I32Tensor = T::arange(0, 1) + pid;
283    let tgt: T::Tensor<i32> = T::load(
284        targets_ptr.add_offsets(tgt_off),
285        None,
286        None,
287        &[],
288        None,
289        None,
290        None,
291        false,
292    );
293
294    // Step 1: store dy * softmax(x) to the full row (uses T::I32Tensor offsets)
295    let dy_bcast = T::broadcast_to(dy, &[BLOCK_SIZE]);
296    let dx_row = dy_bcast * sm;
297    T::store(
298        dx_ptr.add_offsets(row_offs),
299        dx_row,
300        Some(in_row),
301        &[],
302        None,
303        None,
304    );
305
306    // Step 2: subtract dy at target position via atomic_add(-dy)
307    // flat index: pid * n_cols + target (T::Tensor<i32> space)
308    let base: T::Tensor<i32> = T::full::<i32>(&[1], row_base);
309    let flat_off: T::Tensor<i32> = base + tgt;
310    let neg_dy = T::full(&[1], -1.0_f32) * dy;
311    T::atomic_add(dx_ptr.add_offsets(flat_off), neg_dy, None, None, None);
312}
313
314// ── MultiLabelSoftMarginLoss ──────────────────────────────────────────────────
315
316/// Multi-label soft-margin loss forward (element-wise).
317///
318/// Identical to `BCEWithLogitsLoss` per element:
319/// ```text
320/// out = max(x, 0) - x*y + log(1 + exp(-|x|))
321/// ```
322///
323/// Grid: `[ceil(n / BLOCK_SIZE), 1, 1]`, block `[128, 1, 1]`.
324#[kernel]
325pub fn multilabel_soft_margin_loss_forward<T: Triton, const BLOCK_SIZE: i32>(
326    input_ptr: T::Pointer<f32>,
327    target_ptr: T::Pointer<f32>,
328    out_ptr: T::Pointer<f32>,
329    n_elements: i32,
330) where
331    T::I32Tensor: types::Tensor<i32, 1>,
332    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
333    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
334{
335    let pid = T::program_id(Axis::X);
336    let block_start = pid * BLOCK_SIZE;
337    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
338    let in_bounds = offsets.lt(n_elements);
339    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
340
341    let inp = T::load(
342        input_ptr.add_offsets(offsets),
343        Some(in_bounds),
344        Some(zeros),
345        &[],
346        None,
347        None,
348        None,
349        false,
350    );
351    let tgt = T::load(
352        target_ptr.add_offsets(offsets),
353        Some(in_bounds),
354        Some(zeros),
355        &[],
356        None,
357        None,
358        None,
359        false,
360    );
361
362    let one = T::full(&[BLOCK_SIZE], 1.0_f32);
363    let neg_one = T::full(&[BLOCK_SIZE], -1.0_f32);
364    // Numerically stable BCE-with-logits:  max(x,0) - x*t + log(1+exp(-|x|))
365    let relu_x = T::maximum(inp, zeros);
366    let neg_abs_x = neg_one * T::abs(inp);
367    let loss = relu_x - inp * tgt + T::log(one + T::exp(neg_abs_x));
368    T::store(
369        out_ptr.add_offsets(offsets),
370        loss,
371        Some(in_bounds),
372        &[],
373        None,
374        None,
375    );
376}
377
378/// Multi-label soft-margin loss backward (element-wise).
379///
380/// `dx = (sigmoid(x) - target) * dy`
381#[kernel]
382pub fn multilabel_soft_margin_loss_backward<T: Triton, const BLOCK_SIZE: i32>(
383    dy_ptr: T::Pointer<f32>,
384    input_ptr: T::Pointer<f32>,
385    target_ptr: T::Pointer<f32>,
386    dx_ptr: T::Pointer<f32>,
387    n_elements: i32,
388) where
389    T::I32Tensor: types::Tensor<i32, 1>,
390    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
391    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
392{
393    let pid = T::program_id(Axis::X);
394    let block_start = pid * BLOCK_SIZE;
395    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
396    let in_bounds = offsets.lt(n_elements);
397    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
398
399    let dy = T::load(
400        dy_ptr.add_offsets(offsets),
401        Some(in_bounds),
402        Some(zeros),
403        &[],
404        None,
405        None,
406        None,
407        false,
408    );
409    let inp = T::load(
410        input_ptr.add_offsets(offsets),
411        Some(in_bounds),
412        Some(zeros),
413        &[],
414        None,
415        None,
416        None,
417        false,
418    );
419    let tgt = T::load(
420        target_ptr.add_offsets(offsets),
421        Some(in_bounds),
422        Some(zeros),
423        &[],
424        None,
425        None,
426        None,
427        false,
428    );
429
430    let one = T::full(&[BLOCK_SIZE], 1.0_f32);
431    let neg_one = T::full(&[BLOCK_SIZE], -1.0_f32);
432    // sigmoid(x) manually to avoid __nv_sigmoidf
433    let sig = one / (one + T::exp(neg_one * inp));
434    let dx = (sig - tgt) * dy;
435    T::store(
436        dx_ptr.add_offsets(offsets),
437        dx,
438        Some(in_bounds),
439        &[],
440        None,
441        None,
442    );
443}
teeny_kernels/nn/loss/nll.rs

teeny_kernels/nn/loss/
nll.rs
