Skip to main content

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;
12
13pub use helper::apply_schema_quant;
14
15/// Output element type chosen by the user at `DecoderBuilder::with_decode_dtype()`.
16///
17/// The whole post-merge pipeline (boxes, scores, mask coefs, protos) is
18/// emitted in this dtype. `F16` saves ~2× memory bandwidth at the cost of
19/// 10-bit mantissa precision — empirically safe for YOLO-family models.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub enum DecodeDtype {
22    /// 32-bit float output (default).
23    #[default]
24    F32,
25    /// 16-bit float output. Halves the memory bandwidth of the decode
26    /// pipeline; on aarch64 it also keeps the NEON FP16 kernels in their
27    /// native width instead of widening.
28    F16,
29}
30
31/// Activation function applied after dequantization on a logical output.
32///
33/// Sourced from the schema's `activation_required` field. Currently only
34/// `Sigmoid` is wired through the per-scale pipeline; future activations
35/// (e.g. `Softmax` on objectness) extend this enum without ripple.
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
37#[allow(dead_code)] // consumed by per-scale pipeline tasks
38pub(crate) enum Activation {
39    #[default]
40    None,
41    Sigmoid,
42}
43
44impl Activation {
45    /// Translate a schema activation to a per_scale Activation.
46    /// Returns Activation::None when the schema declares no activation.
47    #[allow(dead_code)] // consumed by per-scale pipeline tasks
48    pub(crate) fn from_schema(s: Option<crate::schema::Activation>) -> Self {
49        match s {
50            Some(crate::schema::Activation::Sigmoid) => Self::Sigmoid,
51            _ => Self::None,
52        }
53    }
54}
55
56pub(crate) use outputs::{DecodedOutputBuffers, DecodedOutputsRef};
57pub(crate) use plan::PerScalePlan;
58
59/// Per-scale decoder for schema-v2 per-scale models. Built once at
60/// `DecoderBuilder::build()` time; consumed per-frame via `run()`.
61#[derive(Debug)]
62#[allow(dead_code)] // Wired by Task 24's Decoder integration.
63pub(crate) struct PerScaleDecoder {
64    pub(crate) plan: PerScalePlan,
65    pub(crate) buffers: DecodedOutputBuffers,
66}
67
68impl PerScaleDecoder {
69    /// Build a decoder from a plan, allocating output buffers.
70    #[allow(dead_code)] // Wired by Task 23's builder.
71    pub(crate) fn new(plan: PerScalePlan) -> Self {
72        let buffers = DecodedOutputBuffers::new(
73            plan.out_dtype,
74            plan.total_anchors,
75            plan.num_classes,
76            plan.num_mask_coefs,
77            plan.proto_nhwc_shape.as_deref(),
78        );
79        Self { plan, buffers }
80    }
81
82    /// Decode one frame's worth of inputs.
83    #[allow(dead_code)] // Wired by Task 24.
84    pub(crate) fn run<'a>(
85        &'a mut self,
86        inputs: &[&edgefirst_tensor::TensorDyn],
87    ) -> crate::DecoderResult<DecodedOutputsRef<'a>> {
88        pipeline::run(&self.plan, &mut self.buffers, inputs)
89    }
90}
91
92/// Owned f32 snapshot of pre-NMS per-scale outputs.
93///
94/// Returned by [`crate::Decoder::_testing_run_per_scale_pre_nms`] and
95/// used by integration tests to compare against fixture intermediates
96/// without the noise of NMS ordering.
97#[doc(hidden)]
98pub struct PreNmsCapture {
99    pub boxes_xywh: ndarray::Array2<f32>,
100    pub scores: ndarray::Array2<f32>,
101    pub mask_coefs: Option<ndarray::Array2<f32>>,
102    pub protos: Option<ndarray::Array4<f32>>,
103}