vision_rs/models/yolo/kernels/
detect_decode.rs

1//! Fused dist2bbox decode Triton kernel for YOLO detection post-processing.
2//!
3//! Converts raw LTRB distance predictions from a detection head into
4//! XYWH world-coordinate boxes, fusing the dist2bbox conversion and
5//! stride-scaling into a single kernel pass.
6//!
7//! Layout:
8//!   - `boxes`:    `[B, 4, A]` — raw LTRB distances
9//!   - `anchor_x`: `[A]`       — anchor centre x per anchor
10//!   - `anchor_y`: `[A]`       — anchor centre y per anchor
11//!   - `strides`:  `[A]`       — stride scale per anchor
12//!   - `out`:      `[B, 4, A]` — decoded XYWH boxes in world coordinates
13//!
14//! Parallelism: one CTA per (batch, BLOCK_A-wide anchor tile).
15//! Grid: `B * cdiv(A, BLOCK_A)` flat CTAs.
16
17#![allow(non_snake_case)]
18
19use std::{any::Any, marker::PhantomData, sync::Arc};
20
21use teeny_core::{
22    device::program::ArgVisitor,
23    dtype::{Float, FloatBytes},
24    graph::{CustomOp, Shape},
25    model::{RawPtr, RuntimeOp},
26};
27use teeny_macros::kernel;
28use teeny_triton::triton::{
29    types::{AddOffsets, Comparison},
30    *,
31};
32
33/// Fused dist2bbox + stride-scale decode: LTRB distances → XYWH world coords.
34///
35/// Grid: `B * cdiv(A, BLOCK_A)` — one CTA per (batch element, anchor tile).
36#[allow(clippy::erasing_op, clippy::identity_op)]
37#[kernel]
38pub fn detect_decode_forward<T: Triton, D: Float, const BLOCK_A: i32>(
39    boxes_ptr: T::Pointer<D>,
40    anchor_x_ptr: T::Pointer<D>,
41    anchor_y_ptr: T::Pointer<D>,
42    strides_ptr: T::Pointer<D>,
43    out_ptr: T::Pointer<D>,
44    _B: i32,
45    A: i32,
46) where
47    T::I32Tensor: types::Tensor<i32, 1>,
48    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
49    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
50{
51    let a_tiles = T::cdiv(A, BLOCK_A);
52    let pid_b   = T::program_id(Axis::X) / a_tiles;
53    let a_tile  = T::program_id(Axis::X) % a_tiles;
54    let a_start = a_tile * BLOCK_A;
55
56    let a_offs = T::arange(0, BLOCK_A) + a_start;
57    let mask   = a_offs.lt(A);
58
59    let zeros = T::zeros::<D>(&[BLOCK_A]);
60
61    let anchor_x = T::load(anchor_x_ptr.add_offsets(a_offs), Some(mask), Some(zeros), &[], None, None, None, false);
62    let anchor_y = T::load(anchor_y_ptr.add_offsets(a_offs), Some(mask), Some(zeros), &[], None, None, None, false);
63    let strides  = T::load(strides_ptr.add_offsets(a_offs),  Some(mask), Some(zeros), &[], None, None, None, false);
64
65    let base = pid_b * 4 * A;
66    let dx1 = T::load(boxes_ptr.add_offsets(a_offs + (base + 0 * A)), Some(mask), Some(zeros), &[], None, None, None, false);
67    let dy1 = T::load(boxes_ptr.add_offsets(a_offs + (base + 1 * A)), Some(mask), Some(zeros), &[], None, None, None, false);
68    let dx2 = T::load(boxes_ptr.add_offsets(a_offs + (base + 2 * A)), Some(mask), Some(zeros), &[], None, None, None, false);
69    let dy2 = T::load(boxes_ptr.add_offsets(a_offs + (base + 3 * A)), Some(mask), Some(zeros), &[], None, None, None, false);
70
71    let x1 = anchor_x - dx1;
72    let x2 = anchor_x + dx2;
73    let y1 = anchor_y - dy1;
74    let y2 = anchor_y + dy2;
75
76    let half    = T::full(&[BLOCK_A], D::from_f64(0.5));
77    let cx = (x1 + x2) * half * strides;
78    let cy = (y1 + y2) * half * strides;
79    let w  = (x2 - x1) * strides;
80    let h  = (y2 - y1) * strides;
81
82    T::store(out_ptr.add_offsets(a_offs + (base + 0 * A)), cx, Some(mask), &[], None, None);
83    T::store(out_ptr.add_offsets(a_offs + (base + 1 * A)), cy, Some(mask), &[], None, None);
84    T::store(out_ptr.add_offsets(a_offs + (base + 2 * A)), w,  Some(mask), &[], None, None);
85    T::store(out_ptr.add_offsets(a_offs + (base + 3 * A)), h,  Some(mask), &[], None, None);
86}
87
88// ---------------------------------------------------------------------------
89// DetectDecodeOp — the CustomOp for graph-level representation
90// ---------------------------------------------------------------------------
91
92/// Graph-level representation of the detect_decode op.
93///
94/// Stores precomputed anchor grid and stride data used to build the
95/// `DetectDecodeRuntimeOp` at lowering time via `CustomOp::lower()`.
96pub struct DetectDecodeOp<D: FloatBytes + Send + Sync + 'static> {
97    /// Anchor point x-coordinates, one per anchor.
98    pub anchor_x: Vec<f32>,
99    /// Anchor point y-coordinates, one per anchor.
100    pub anchor_y: Vec<f32>,
101    /// Per-anchor stride (8/16/32 for the respective FPN scale).
102    pub strides: Vec<f32>,
103    /// Kernel launch block size along the anchor dimension.
104    pub block_a: i32,
105    _phantom: PhantomData<D>,
106}
107
108impl<D: FloatBytes + Send + Sync + 'static> DetectDecodeOp<D> {
109    /// Creates the op with a precomputed anchor grid and launch block size.
110    pub fn new(anchor_x: Vec<f32>, anchor_y: Vec<f32>, strides: Vec<f32>, block_a: i32) -> Self {
111        Self { anchor_x, anchor_y, strides, block_a, _phantom: PhantomData }
112    }
113}
114
115impl<D: FloatBytes + Send + Sync + 'static> CustomOp for DetectDecodeOp<D> {
116    fn name(&self) -> &str { "yolo.detect_decode" }
117
118    fn infer_output_shape(&self, input_shapes: &[&Shape]) -> Shape {
119        // boxes [B, 4, A] → [B, 4, A]: shape-preserving
120        input_shapes[0].clone()
121    }
122
123    fn as_any(&self) -> &dyn Any { self }
124
125    fn lower(&self) -> Option<(String, String, String, Arc<dyn RuntimeOp>)> {
126        let kernel = DetectDecodeForward::<D>::new(self.block_a);
127        let runtime_op: Arc<dyn RuntimeOp> = Arc::new(DetectDecodeRuntimeOp::<D>::new(
128            self.anchor_x.clone(),
129            self.anchor_y.clone(),
130            self.strides.clone(),
131            self.block_a,
132        ));
133        Some(("detect_decode_forward".to_string(), kernel.source, "entry_point".to_string(), runtime_op))
134    }
135}
136
137// ---------------------------------------------------------------------------
138// DetectDecodeRuntimeOp — the RuntimeOp for kernel dispatch
139// ---------------------------------------------------------------------------
140
141/// Runtime dispatch for the detect_decode kernel.
142///
143/// Anchor grid and strides are stored here so they can be uploaded to device
144/// parameter buffers via [`RuntimeOp::param_init_data`] at model-load time.
145pub struct DetectDecodeRuntimeOp<D: FloatBytes + Send + Sync + 'static> {
146    anchor_x: Vec<f32>,
147    anchor_y: Vec<f32>,
148    strides: Vec<f32>,
149    block_a: i32,
150    _phantom: PhantomData<D>,
151}
152
153impl<D: FloatBytes + Send + Sync + 'static> DetectDecodeRuntimeOp<D> {
154    /// Creates the runtime op with a precomputed anchor grid and launch block size.
155    pub fn new(
156        anchor_x: Vec<f32>,
157        anchor_y: Vec<f32>,
158        strides: Vec<f32>,
159        block_a: i32,
160    ) -> Self {
161        Self { anchor_x, anchor_y, strides, block_a, _phantom: PhantomData }
162    }
163}
164
165impl<D: FloatBytes + Send + Sync + 'static> RuntimeOp for DetectDecodeRuntimeOp<D> {
166    fn n_activation_inputs(&self) -> usize { 1 }
167
168    fn param_shapes(
169        &self,
170        input_shapes: &[&[usize]],
171        _output_shape: &[usize],
172    ) -> Vec<Vec<usize>> {
173        // input_shapes[0] is boxes: [B, 4, A]
174        let a = input_shapes[0][2];
175        vec![vec![a], vec![a], vec![a]] // anchor_x, anchor_y, strides
176    }
177
178    fn param_init_data(&self, param_idx: usize) -> Option<Vec<u8>> {
179        let data: &[f32] = match param_idx {
180            0 => &self.anchor_x,
181            1 => &self.anchor_y,
182            2 => &self.strides,
183            _ => return None,
184        };
185        // Upload the anchor grid / strides in the device buffer's element type
186        // `D`, converting from the host-side f32 geometry generically.
187        Some(data.iter().flat_map(|&f| D::from_f64(f as f64).to_le_bytes()).collect())
188    }
189
190    fn pack_args(
191        &self,
192        inputs: &[(RawPtr, &[usize])],
193        params: &[RawPtr],
194        output: RawPtr,
195        output_shape: &[usize],
196        _output_row_stride: i32,
197        visitor: &mut dyn ArgVisitor,
198    ) {
199        let b = output_shape[0] as i32;
200        let a = output_shape[2] as i32;
201        visitor.visit_ptr(inputs[0].0); // boxes_ptr
202        visitor.visit_ptr(params[0]);   // anchor_x_ptr
203        visitor.visit_ptr(params[1]);   // anchor_y_ptr
204        visitor.visit_ptr(params[2]);   // strides_ptr
205        visitor.visit_ptr(output);      // out_ptr
206        visitor.visit_i32(b);           // _B
207        visitor.visit_i32(a);           // A
208    }
209
210    fn block(&self) -> [u32; 3] { [self.block_a as u32, 1, 1] }
211
212    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
213        let b = output_shape[0];
214        let a = output_shape[2];
215        let a_tiles = a.div_ceil(self.block_a as usize);
216        [(b * a_tiles) as u32, 1, 1]
217    }
218}
vision_rs/models/yolo/kernels/detect_decode.rs

vision_rs/models/yolo/kernels/
detect_decode.rs
