teeny_kernels/nn/loss/
ranking.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// ── MarginRankingLoss ─────────────────────────────────────────────────────────
26
27/// Margin ranking loss forward (element-wise).
28///
29/// `out[i] = max(0, -y[i] * (x1[i] - x2[i]) + margin)`
30///
31/// Grid: `[ceil(n / BLOCK_SIZE), 1, 1]`
32#[kernel]
33pub fn margin_ranking_loss_forward<T: Triton, const BLOCK_SIZE: i32>(
34    x1_ptr: T::Pointer<f32>,
35    x2_ptr: T::Pointer<f32>,
36    y_ptr: T::Pointer<f32>,
37    out_ptr: T::Pointer<f32>,
38    n_elements: i32,
39    margin: f32,
40) where
41    T::I32Tensor: types::Tensor<i32, 1>,
42    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
43    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
44{
45    let pid = T::program_id(Axis::X);
46    let block_start = pid * BLOCK_SIZE;
47    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
48    let in_bounds = offsets.lt(n_elements);
49    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
50
51    let x1 = T::load(
52        x1_ptr.add_offsets(offsets),
53        Some(in_bounds),
54        Some(zeros),
55        &[],
56        None,
57        None,
58        None,
59        false,
60    );
61    let x2 = T::load(
62        x2_ptr.add_offsets(offsets),
63        Some(in_bounds),
64        Some(zeros),
65        &[],
66        None,
67        None,
68        None,
69        false,
70    );
71    let y = T::load(
72        y_ptr.add_offsets(offsets),
73        Some(in_bounds),
74        Some(zeros),
75        &[],
76        None,
77        None,
78        None,
79        false,
80    );
81
82    let margin_t = T::full(&[BLOCK_SIZE], margin);
83    let neg_one = T::full(&[BLOCK_SIZE], -1.0_f32);
84    // hinge = max(0, -y*(x1-x2) + margin)
85    let hinge = T::maximum(neg_one * y * (x1 - x2) + margin_t, zeros);
86    T::store(
87        out_ptr.add_offsets(offsets),
88        hinge,
89        Some(in_bounds),
90        &[],
91        None,
92        None,
93    );
94}
95
96/// Margin ranking loss backward (element-wise).
97///
98/// `dx1[i] = -y[i] * dy[i]` if hinge > 0, else 0.
99/// `dx2[i] =  y[i] * dy[i]` if hinge > 0, else 0.
100#[kernel]
101pub fn margin_ranking_loss_backward<T: Triton, const BLOCK_SIZE: i32>(
102    dy_ptr: T::Pointer<f32>,
103    x1_ptr: T::Pointer<f32>,
104    x2_ptr: T::Pointer<f32>,
105    y_ptr: T::Pointer<f32>,
106    dx1_ptr: T::Pointer<f32>,
107    dx2_ptr: T::Pointer<f32>,
108    n_elements: i32,
109    margin: f32,
110) where
111    T::I32Tensor: types::Tensor<i32, 1>,
112    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
113    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
114{
115    let pid = T::program_id(Axis::X);
116    let block_start = pid * BLOCK_SIZE;
117    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
118    let in_bounds = offsets.lt(n_elements);
119    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
120
121    let dy = T::load(
122        dy_ptr.add_offsets(offsets),
123        Some(in_bounds),
124        Some(zeros),
125        &[],
126        None,
127        None,
128        None,
129        false,
130    );
131    let x1 = T::load(
132        x1_ptr.add_offsets(offsets),
133        Some(in_bounds),
134        Some(zeros),
135        &[],
136        None,
137        None,
138        None,
139        false,
140    );
141    let x2 = T::load(
142        x2_ptr.add_offsets(offsets),
143        Some(in_bounds),
144        Some(zeros),
145        &[],
146        None,
147        None,
148        None,
149        false,
150    );
151    let y = T::load(
152        y_ptr.add_offsets(offsets),
153        Some(in_bounds),
154        Some(zeros),
155        &[],
156        None,
157        None,
158        None,
159        false,
160    );
161
162    let margin_t = T::full(&[BLOCK_SIZE], margin);
163    let neg_one = T::full(&[BLOCK_SIZE], -1.0_f32);
164    let pre_hinge = neg_one * y * (x1 - x2) + margin_t;
165    let active = T::gt(pre_hinge, zeros);
166
167    let dx1 = T::where_(active, neg_one * y * dy, zeros);
168    let dx2 = T::where_(active, y * dy, zeros);
169    T::store(
170        dx1_ptr.add_offsets(offsets),
171        dx1,
172        Some(in_bounds),
173        &[],
174        None,
175        None,
176    );
177    T::store(
178        dx2_ptr.add_offsets(offsets),
179        dx2,
180        Some(in_bounds),
181        &[],
182        None,
183        None,
184    );
185}
186
187// ── HingeEmbeddingLoss ────────────────────────────────────────────────────────
188
189/// Hinge embedding loss forward (element-wise).
190///
191/// ```text
192/// out[i] = x[i]                     if y[i] ==  1
193///        = max(0, margin - x[i])    if y[i] == -1
194/// ```
195///
196/// Grid: `[ceil(n / BLOCK_SIZE), 1, 1]`
197#[kernel]
198pub fn hinge_embedding_loss_forward<T: Triton, const BLOCK_SIZE: i32>(
199    inp_ptr: T::Pointer<f32>,
200    y_ptr: T::Pointer<f32>,
201    out_ptr: T::Pointer<f32>,
202    n_elements: i32,
203    margin: f32,
204) where
205    T::I32Tensor: types::Tensor<i32, 1>,
206    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
207    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
208{
209    let pid = T::program_id(Axis::X);
210    let block_start = pid * BLOCK_SIZE;
211    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
212    let in_bounds = offsets.lt(n_elements);
213    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
214
215    let inp = T::load(
216        inp_ptr.add_offsets(offsets),
217        Some(in_bounds),
218        Some(zeros),
219        &[],
220        None,
221        None,
222        None,
223        false,
224    );
225    let y = T::load(
226        y_ptr.add_offsets(offsets),
227        Some(in_bounds),
228        Some(zeros),
229        &[],
230        None,
231        None,
232        None,
233        false,
234    );
235
236    let margin_t = T::full(&[BLOCK_SIZE], margin);
237    // y > 0 means y == 1
238    let y_is_pos = T::gt(y, zeros);
239    let hinge = T::maximum(margin_t - inp, zeros);
240    let out = T::where_(y_is_pos, inp, hinge);
241    T::store(
242        out_ptr.add_offsets(offsets),
243        out,
244        Some(in_bounds),
245        &[],
246        None,
247        None,
248    );
249}
250
251/// Hinge embedding loss backward (element-wise).
252///
253/// ```text
254/// dx[i] = dy[i]     if y[i] ==  1
255///       = -dy[i]    if y[i] == -1 and margin - x[i] > 0
256///       = 0         otherwise
257/// ```
258#[kernel]
259pub fn hinge_embedding_loss_backward<T: Triton, const BLOCK_SIZE: i32>(
260    dy_ptr: T::Pointer<f32>,
261    inp_ptr: T::Pointer<f32>,
262    y_ptr: T::Pointer<f32>,
263    dx_ptr: T::Pointer<f32>,
264    n_elements: i32,
265    margin: f32,
266) where
267    T::I32Tensor: types::Tensor<i32, 1>,
268    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
269    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
270{
271    let pid = T::program_id(Axis::X);
272    let block_start = pid * BLOCK_SIZE;
273    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
274    let in_bounds = offsets.lt(n_elements);
275    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
276
277    let dy = T::load(
278        dy_ptr.add_offsets(offsets),
279        Some(in_bounds),
280        Some(zeros),
281        &[],
282        None,
283        None,
284        None,
285        false,
286    );
287    let inp = T::load(
288        inp_ptr.add_offsets(offsets),
289        Some(in_bounds),
290        Some(zeros),
291        &[],
292        None,
293        None,
294        None,
295        false,
296    );
297    let y = T::load(
298        y_ptr.add_offsets(offsets),
299        Some(in_bounds),
300        Some(zeros),
301        &[],
302        None,
303        None,
304        None,
305        false,
306    );
307
308    let margin_t = T::full(&[BLOCK_SIZE], margin);
309    let neg_one = T::full(&[BLOCK_SIZE], -1.0_f32);
310    let y_is_pos = T::gt(y, zeros);
311    // Active for y == -1: margin - x > 0
312    let neg_active = T::gt(margin_t - inp, zeros);
313    let dx_neg = T::where_(neg_active, neg_one * dy, zeros);
314    let dx = T::where_(y_is_pos, dy, dx_neg);
315    T::store(
316        dx_ptr.add_offsets(offsets),
317        dx,
318        Some(in_bounds),
319        &[],
320        None,
321        None,
322    );
323}
324
325// ── MultiMarginLoss ───────────────────────────────────────────────────────────
326
327/// Multi-margin loss forward (per-row).
328///
329/// `out[n] = (1/n_cols) * sum_{j != target[n]} max(0, margin - x[n,target] + x[n,j])`
330///
331/// Computed as:
332/// `(sum_all max(0, margin - x_t + x_j) - max(0, margin)) / n_cols`
333///
334/// Grid: `[n_rows, 1, 1]` — one CTA per row.
335/// `BLOCK_SIZE` must equal `next_power_of_two(n_cols)`.
336#[kernel]
337pub fn multi_margin_loss_forward<T: Triton, const BLOCK_SIZE: i32>(
338    input_ptr: T::Pointer<f32>,
339    targets_ptr: T::Pointer<i32>,
340    out_ptr: T::Pointer<f32>,
341    _n_rows: i32,
342    n_cols: i32,
343    margin: f32,
344) where
345    T::I32Tensor: types::Tensor<i32, 1>,
346    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
347    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
348    T::Pointer<i32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<i32>>>,
349    T::Tensor<i32>: types::Tensor<i32, 1>,
350    T::Pointer<f32>: AddOffsets<i32, 1, T::Tensor<i32>, Output = T::Tensor<T::Pointer<f32>>>,
351{
352    let pid = T::program_id(Axis::X);
353    let row_base = pid * n_cols;
354    let col_offs: T::I32Tensor = T::arange(0, BLOCK_SIZE);
355    let row_offs: T::I32Tensor = col_offs + row_base;
356    let in_row = col_offs.lt(n_cols);
357    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
358
359    let row = T::load(
360        input_ptr.add_offsets(row_offs),
361        Some(in_row),
362        Some(zeros),
363        &[],
364        None,
365        None,
366        None,
367        false,
368    );
369
370    // Load target index
371    let tgt_off: T::I32Tensor = T::arange(0, 1) + pid;
372    let tgt: T::Tensor<i32> = T::load(
373        targets_ptr.add_offsets(tgt_off),
374        None,
375        None,
376        &[],
377        None,
378        None,
379        None,
380        false,
381    );
382
383    // Load x[target] using flat index
384    let base: T::Tensor<i32> = T::full::<i32>(&[1], row_base);
385    let flat_off: T::Tensor<i32> = base + tgt;
386    let x_t: T::Tensor<f32> = T::load(
387        input_ptr.add_offsets(flat_off),
388        None,
389        None,
390        &[],
391        None,
392        None,
393        None,
394        false,
395    );
396
397    // hinge for all j (including t): max(0, margin - x_t + x_j)
398    let margin_t = T::full(&[BLOCK_SIZE], margin);
399    let x_t_bcast = T::broadcast_to(x_t, &[BLOCK_SIZE]);
400    let hinge_all = T::maximum(margin_t - x_t_bcast + row, zeros);
401
402    // sum all hinges, subtract the target's own contribution: max(0, margin)
403    let sum_all = T::sum(hinge_all, Some(0), true); // shape [1]
404    // max(0, margin) as a tensor to avoid scalar f32 comparison in kernel context
405    let tgt_contrib = T::maximum(T::full::<f32>(&[1], margin), T::zeros::<f32>(&[1]));
406    // Cast n_cols to f32 via tensor cast (scalar `as f32` produces ub.poison in kernels)
407    let n_cols_f = T::cast::<i32, f32>(T::full::<i32>(&[1], n_cols), None, false);
408    let loss = (sum_all - tgt_contrib) / n_cols_f;
409
410    let out_off: T::I32Tensor = T::arange(0, 1) + pid;
411    T::store(out_ptr.add_offsets(out_off), loss, None, &[], None, None);
412}
413
414/// Multi-margin loss backward (per-row).
415///
416/// Grid: `[n_rows, 1, 1]`.  `BLOCK_SIZE` must equal `next_power_of_two(n_cols)`.
417///
418/// Two-step write:
419/// 1. Store `(dy / n_cols) * active_j` for all j in the row.
420/// 2. Atomic-add correction at target: `-(dy / n_cols) * sum(active_all)`.
421#[kernel]
422pub fn multi_margin_loss_backward<T: Triton, const BLOCK_SIZE: i32>(
423    dy_ptr: T::Pointer<f32>,
424    input_ptr: T::Pointer<f32>,
425    targets_ptr: T::Pointer<i32>,
426    dx_ptr: T::Pointer<f32>,
427    _n_rows: i32,
428    n_cols: i32,
429    margin: f32,
430) where
431    T::I32Tensor: types::Tensor<i32, 1>,
432    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
433    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
434    T::Pointer<i32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<i32>>>,
435    T::Tensor<i32>: types::Tensor<i32, 1>,
436    T::Pointer<f32>: AddOffsets<i32, 1, T::Tensor<i32>, Output = T::Tensor<T::Pointer<f32>>>,
437{
438    let pid = T::program_id(Axis::X);
439    let row_base = pid * n_cols;
440    let col_offs: T::I32Tensor = T::arange(0, BLOCK_SIZE);
441    let row_offs: T::I32Tensor = col_offs + row_base;
442    let in_row = col_offs.lt(n_cols);
443    let zeros = T::zeros::<f32>(&[BLOCK_SIZE]);
444
445    // Load upstream gradient
446    let dy_off: T::I32Tensor = T::arange(0, 1) + pid;
447    let dy: T::Tensor<f32> = T::load(
448        dy_ptr.add_offsets(dy_off),
449        None,
450        None,
451        &[],
452        None,
453        None,
454        None,
455        false,
456    );
457
458    let row = T::load(
459        input_ptr.add_offsets(row_offs),
460        Some(in_row),
461        Some(zeros),
462        &[],
463        None,
464        None,
465        None,
466        false,
467    );
468
469    // Load target index
470    let tgt_off: T::I32Tensor = T::arange(0, 1) + pid;
471    let tgt: T::Tensor<i32> = T::load(
472        targets_ptr.add_offsets(tgt_off),
473        None,
474        None,
475        &[],
476        None,
477        None,
478        None,
479        false,
480    );
481
482    // Load x[target] using flat index
483    let base: T::Tensor<i32> = T::full::<i32>(&[1], row_base);
484    let flat_off: T::Tensor<i32> = base + tgt;
485    let x_t: T::Tensor<f32> = T::load(
486        input_ptr.add_offsets(flat_off),
487        None,
488        None,
489        &[],
490        None,
491        None,
492        None,
493        false,
494    );
495
496    let margin_t = T::full(&[BLOCK_SIZE], margin);
497    let x_t_bcast = T::broadcast_to(x_t, &[BLOCK_SIZE]);
498    let ones = T::full(&[BLOCK_SIZE], 1.0_f32);
499
500    // active[j] = (margin - x_t + x_j > 0)
501    let active = T::gt(margin_t - x_t_bcast + row, zeros);
502    let active_f = T::where_(active, ones, zeros);
503
504    // sum of all active (including target position)
505    let sum_active = T::sum(active_f, Some(0), true); // shape [1]
506
507    let n_cols_f = T::cast::<i32, f32>(T::full::<i32>(&[1], n_cols), None, false);
508    let dy_over_n = dy / n_cols_f;
509
510    // Step 1: store (dy/n_cols) * active_f for the whole row
511    let dy_bcast = T::broadcast_to(dy_over_n, &[BLOCK_SIZE]);
512    let dx_row = dy_bcast * active_f;
513    T::store(
514        dx_ptr.add_offsets(row_offs),
515        dx_row,
516        Some(in_row),
517        &[],
518        None,
519        None,
520    );
521
522    // Step 2: atomic_add correction at target to fix target position
523    // After step 1, position t has (dy/n_cols)*active_t; correct value is -(dy/n_cols)*count_not_t
524    // correction = -(dy/n_cols)*sum_active_all
525    let neg_one = T::full(&[1], -1.0_f32);
526    let correction = neg_one * dy_over_n * sum_active;
527    T::atomic_add(dx_ptr.add_offsets(flat_off), correction, None, None, None);
528}
teeny_kernels/nn/loss/ranking.rs

teeny_kernels/nn/loss/
ranking.rs
