edgefirst_decoder/per_scale/
mod.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Per-scale quantized decoder — see
5//! `.claude/plans/2026-04-28-per-scale-decoder-optimized-design.md`.
6
7pub mod helper;
8pub(crate) mod kernels;
9pub(crate) mod outputs;
10pub(crate) mod pipeline;
11pub(crate) mod plan;
12pub(crate) mod prefilter;
13
14pub use helper::apply_schema_quant;
15
16/// Output element type chosen by the user at `DecoderBuilder::with_decode_dtype()`.
17///
18/// The whole post-merge pipeline (boxes, scores, mask coefs, protos) is
19/// emitted in this dtype. `F16` saves ~2× memory bandwidth at the cost of
20/// 10-bit mantissa precision — empirically safe for YOLO-family models.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
22pub enum DecodeDtype {
23    /// 32-bit float output (default).
24    #[default]
25    F32,
26    /// 16-bit float output. Halves the memory bandwidth of the decode
27    /// pipeline; on aarch64 it also keeps the NEON FP16 kernels in their
28    /// native width instead of widening.
29    F16,
30}
31
32/// Activation function applied after dequantization on a logical output.
33///
34/// Sourced from the schema's `activation_required` field. Currently only
35/// `Sigmoid` is wired through the per-scale pipeline; future activations
36/// (e.g. `Softmax` on objectness) extend this enum without ripple.
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
38#[allow(dead_code)] // consumed by per-scale pipeline tasks
39pub(crate) enum Activation {
40    #[default]
41    None,
42    Sigmoid,
43}
44
45impl Activation {
46    /// Translate a schema activation to a per_scale Activation.
47    /// Returns Activation::None when the schema declares no activation.
48    #[allow(dead_code)] // consumed by per-scale pipeline tasks
49    pub(crate) fn from_schema(s: Option<crate::schema::Activation>) -> Self {
50        match s {
51            Some(crate::schema::Activation::Sigmoid) => Self::Sigmoid,
52            _ => Self::None,
53        }
54    }
55}
56
57pub(crate) use outputs::{DecodedOutputBuffers, DecodedOutputsRef};
58pub(crate) use plan::PerScalePlan;
59
60/// Per-scale decoder for schema-v2 per-scale models. Built once at
61/// `DecoderBuilder::build()` time; consumed per-frame via `run()`.
62#[derive(Debug)]
63#[allow(dead_code)] // Wired by Task 24's Decoder integration.
64pub(crate) struct PerScaleDecoder {
65    pub(crate) plan: PerScalePlan,
66    pub(crate) buffers: DecodedOutputBuffers,
67}
68
69impl PerScaleDecoder {
70    /// Build a decoder from a plan, allocating output buffers.
71    #[allow(dead_code)] // Wired by Task 23's builder.
72    pub(crate) fn new(plan: PerScalePlan) -> Self {
73        let buffers = DecodedOutputBuffers::new(
74            plan.out_dtype,
75            plan.total_anchors,
76            plan.num_classes,
77            plan.num_mask_coefs,
78            plan.proto_nhwc_shape.as_deref(),
79        );
80        Self { plan, buffers }
81    }
82
83    /// Decode one frame's worth of inputs.
84    ///
85    /// With a `score_threshold`, only anchors that can reach it are
86    /// decoded and the outputs hold just those, in anchor order; `None`
87    /// decodes every anchor.
88    #[allow(dead_code)] // Wired by Task 24.
89    pub(crate) fn run<'a>(
90        &'a mut self,
91        inputs: &[&edgefirst_tensor::TensorDyn],
92        score_threshold: Option<f32>,
93    ) -> crate::DecoderResult<DecodedOutputsRef<'a>> {
94        pipeline::run(&self.plan, &mut self.buffers, inputs, score_threshold)
95    }
96}
97
98/// Owned f32 snapshot of pre-NMS per-scale outputs.
99///
100/// Returned by [`crate::Decoder::_testing_run_per_scale_pre_nms`] and
101/// used by integration tests to compare against fixture intermediates
102/// without the noise of NMS ordering.
103#[doc(hidden)]
104pub struct PreNmsCapture {
105    pub boxes_xywh: ndarray::Array2<f32>,
106    pub scores: ndarray::Array2<f32>,
107    pub mask_coefs: Option<ndarray::Array2<f32>>,
108    pub protos: Option<ndarray::Array4<f32>>,
109}
110
111/// Kernel-level hooks for the decoder benchmark. Not a stable API.
112#[doc(hidden)]
113pub mod bench {
114    use super::kernels::dfl_lut::{decode_dfl_lut, DflExpTable};
115    use super::kernels::grids::make_anchor_grid;
116    use super::kernels::level_box;
117    use super::plan::{Layout, LevelPlan};
118    use crate::Quantization;
119
120    /// A DFL box kernel for 8-bit bins.
121    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
122    pub enum DflKernel {
123        /// Dequantize, then f32 libm softmax per side (portable).
124        Softmax,
125        /// Exp table, portable.
126        Lut,
127        /// Dequantize, then fused NEON f32 softmax and weighted sum.
128        NeonFused,
129        /// Dequantize, then fused NEON softmax with an FP16 exp.
130        NeonFusedFp16,
131        /// Exp table with the four sides in NEON lanes.
132        LutNeon,
133    }
134
135    impl DflKernel {
136        /// Every kernel this CPU can run.
137        pub fn available() -> Vec<Self> {
138            let v = vec![Self::Softmax, Self::Lut];
139            #[cfg(target_arch = "aarch64")]
140            let v = {
141                let mut v = v;
142                v.extend([Self::NeonFused, Self::LutNeon]);
143                if std::arch::is_aarch64_feature_detected!("fp16") {
144                    v.push(Self::NeonFusedFp16);
145                }
146                v
147            };
148            v
149        }
150    }
151
152    /// One DFL level (`reg_max` 16) and its quantization.
153    pub struct DflLevel {
154        plan: LevelPlan,
155        q: Quantization,
156        table: DflExpTable,
157    }
158
159    impl DflLevel {
160        pub fn new(w: usize, h: usize, stride: f32, scale: f32, zero_point: i32) -> Self {
161            let (grid_x, grid_y) = make_anchor_grid(h, w);
162            Self {
163                plan: LevelPlan {
164                    stride,
165                    h,
166                    w,
167                    reg_max: 16,
168                    anchor_offset: 0,
169                    grid_x,
170                    grid_y,
171                    box_shape: vec![1, h, w, 64].into_boxed_slice(),
172                    score_shape: vec![1, h, w, 1].into_boxed_slice(),
173                    mc_shape: None,
174                    layout: Layout::Nhwc,
175                },
176                q: Quantization::new(scale, zero_point),
177                table: DflExpTable::new(scale),
178            }
179        }
180
181        pub fn anchors(&self) -> usize {
182            self.plan.h * self.plan.w
183        }
184
185        /// Decode `input` (`anchors() * 64` bins) into `dst` (`anchors() * 4`).
186        pub fn decode(&self, kernel: DflKernel, input: &[i8], dst: &mut [f32]) {
187            let (p, q) = (&self.plan, self.q);
188            match kernel {
189                DflKernel::Softmax => level_box::decode_box_level_dfl_i8_to_f32(input, q, p, dst),
190                DflKernel::Lut => decode_dfl_lut(input, &self.table, p, dst),
191                #[cfg(target_arch = "aarch64")]
192                DflKernel::NeonFused => {
193                    level_box::decode_box_level_dfl_i8_to_f32_neon_fused(input, q, p, dst)
194                }
195                #[cfg(target_arch = "aarch64")]
196                DflKernel::NeonFusedFp16 => {
197                    assert!(std::arch::is_aarch64_feature_detected!("fp16"));
198                    level_box::decode_box_level_dfl_i8_to_f32_neon_fused_fp16(input, q, p, dst)
199                }
200                #[cfg(target_arch = "aarch64")]
201                DflKernel::LutNeon => {
202                    super::kernels::dfl_lut::decode_dfl_lut_neon(input, &self.table, p, dst)
203                }
204                #[cfg(not(target_arch = "aarch64"))]
205                _ => panic!("{kernel:?} is not available on this CPU"),
206            }
207        }
208    }
209}
edgefirst_decoder/per_scale/mod.rs

edgefirst_decoder/per_scale/
mod.rs
