Skip to main content

edgefirst_image/cpu/
masks.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use super::CPUProcessor;
5use crate::Result;
6use edgefirst_decoder::{DetectBox, Segmentation};
7use ndarray::Axis;
8use rayon::prelude::*;
9
10impl CPUProcessor {
11    #[allow(clippy::too_many_arguments)]
12    pub(super) fn render_modelpack_segmentation(
13        &mut self,
14        dst_w: usize,
15        dst_h: usize,
16        dst_rs: usize,
17        dst_c: usize,
18        dst_slice: &mut [u8],
19        segmentation: &Segmentation,
20        opacity: f32,
21    ) -> Result<()> {
22        use ndarray_stats::QuantileExt;
23
24        let seg = &segmentation.segmentation;
25        let [seg_height, seg_width, seg_classes] = *seg.shape() else {
26            unreachable!("Array3 did not have [usize; 3] as shape");
27        };
28        let start_y = (dst_h as f32 * segmentation.ymin).round();
29        let end_y = (dst_h as f32 * segmentation.ymax).round();
30        let start_x = (dst_w as f32 * segmentation.xmin).round();
31        let end_x = (dst_w as f32 * segmentation.xmax).round();
32
33        let scale_x = (seg_width as f32 - 1.0) / ((end_x - start_x) - 1.0);
34        let scale_y = (seg_height as f32 - 1.0) / ((end_y - start_y) - 1.0);
35
36        let start_x_u = (start_x as usize).min(dst_w);
37        let start_y_u = (start_y as usize).min(dst_h);
38        let end_x_u = (end_x as usize).min(dst_w);
39        let end_y_u = (end_y as usize).min(dst_h);
40
41        let argmax = seg.map_axis(Axis(2), |r| r.argmax().unwrap());
42        let get_value_at_nearest = |x: f32, y: f32| -> usize {
43            let x = x.round() as usize;
44            let y = y.round() as usize;
45            argmax
46                .get([y.min(seg_height - 1), x.min(seg_width - 1)])
47                .copied()
48                .unwrap_or(0)
49        };
50
51        for y in start_y_u..end_y_u {
52            for x in start_x_u..end_x_u {
53                let seg_x = (x as f32 - start_x) * scale_x;
54                let seg_y = (y as f32 - start_y) * scale_y;
55                let label = get_value_at_nearest(seg_x, seg_y);
56
57                if label == seg_classes - 1 {
58                    continue;
59                }
60
61                let color = self.colors[label % self.colors.len()];
62
63                let alpha = if opacity == 1.0 {
64                    color[3] as u16
65                } else {
66                    (color[3] as f32 * opacity).round() as u16
67                };
68
69                let dst_index = (y * dst_rs) + (x * dst_c);
70                for c in 0..3 {
71                    dst_slice[dst_index + c] = ((color[c] as u16 * alpha
72                        + dst_slice[dst_index + c] as u16 * (255 - alpha))
73                        / 255) as u8;
74                }
75            }
76        }
77
78        Ok(())
79    }
80
81    #[allow(clippy::too_many_arguments)]
82    pub(super) fn render_yolo_segmentation(
83        &mut self,
84        dst_w: usize,
85        dst_h: usize,
86        dst_rs: usize,
87        dst_c: usize,
88        dst_slice: &mut [u8],
89        segmentation: &Segmentation,
90        class: usize,
91        opacity: f32,
92    ) -> Result<()> {
93        let seg = &segmentation.segmentation;
94        let [seg_height, seg_width, classes] = *seg.shape() else {
95            unreachable!("Array3 did not have [usize;3] as shape");
96        };
97        debug_assert_eq!(classes, 1);
98
99        let start_y = (dst_h as f32 * segmentation.ymin).round();
100        let end_y = (dst_h as f32 * segmentation.ymax).round();
101        let start_x = (dst_w as f32 * segmentation.xmin).round();
102        let end_x = (dst_w as f32 * segmentation.xmax).round();
103
104        let scale_x = (seg_width as f32 - 1.0) / ((end_x - start_x) - 1.0);
105        let scale_y = (seg_height as f32 - 1.0) / ((end_y - start_y) - 1.0);
106
107        let start_x_u = (start_x as usize).min(dst_w);
108        let start_y_u = (start_y as usize).min(dst_h);
109        let end_x_u = (end_x as usize).min(dst_w);
110        let end_y_u = (end_y as usize).min(dst_h);
111
112        for y in start_y_u..end_y_u {
113            for x in start_x_u..end_x_u {
114                let seg_x = ((x as f32 - start_x) * scale_x) as usize;
115                let seg_y = ((y as f32 - start_y) * scale_y) as usize;
116                let val = *seg.get([seg_y, seg_x, 0]).unwrap_or(&0);
117
118                if val < 127 {
119                    continue;
120                }
121
122                let color = self.colors[class % self.colors.len()];
123
124                let alpha = if opacity == 1.0 {
125                    color[3] as u16
126                } else {
127                    (color[3] as f32 * opacity).round() as u16
128                };
129
130                let dst_index = (y * dst_rs) + (x * dst_c);
131                for c in 0..3 {
132                    dst_slice[dst_index + c] = ((color[c] as u16 * alpha
133                        + dst_slice[dst_index + c] as u16 * (255 - alpha))
134                        / 255) as u8;
135                }
136            }
137        }
138
139        Ok(())
140    }
141
142    #[allow(clippy::too_many_arguments)]
143    pub(super) fn render_box(
144        &mut self,
145        dst_w: usize,
146        dst_h: usize,
147        dst_rs: usize,
148        dst_c: usize,
149        dst_slice: &mut [u8],
150        detect: &[DetectBox],
151        color_mode: crate::ColorMode,
152    ) -> Result<()> {
153        const LINE_THICKNESS: usize = 3;
154
155        for (idx, d) in detect.iter().enumerate() {
156            use edgefirst_decoder::BoundingBox;
157
158            let color_index = color_mode.index(idx, d.label);
159            let [r, g, b, _] = self.colors[color_index % self.colors.len()];
160            let bbox = d.bbox.to_canonical();
161            let bbox = BoundingBox {
162                xmin: bbox.xmin.clamp(0.0, 1.0),
163                ymin: bbox.ymin.clamp(0.0, 1.0),
164                xmax: bbox.xmax.clamp(0.0, 1.0),
165                ymax: bbox.ymax.clamp(0.0, 1.0),
166            };
167            let inner = [
168                ((dst_w - 1) as f32 * bbox.xmin - 0.5).round() as usize,
169                ((dst_h - 1) as f32 * bbox.ymin - 0.5).round() as usize,
170                ((dst_w - 1) as f32 * bbox.xmax + 0.5).round() as usize,
171                ((dst_h - 1) as f32 * bbox.ymax + 0.5).round() as usize,
172            ];
173
174            let outer = [
175                inner[0].saturating_sub(LINE_THICKNESS),
176                inner[1].saturating_sub(LINE_THICKNESS),
177                (inner[2] + LINE_THICKNESS).min(dst_w),
178                (inner[3] + LINE_THICKNESS).min(dst_h),
179            ];
180
181            // top line
182            for y in outer[1] + 1..=inner[1] {
183                for x in outer[0] + 1..outer[2] {
184                    let index = (y * dst_rs) + (x * dst_c);
185                    dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
186                }
187            }
188
189            // left and right lines
190            for y in inner[1]..inner[3] {
191                for x in outer[0] + 1..=inner[0] {
192                    let index = (y * dst_rs) + (x * dst_c);
193                    dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
194                }
195
196                for x in inner[2]..outer[2] {
197                    let index = (y * dst_rs) + (x * dst_c);
198                    dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
199                }
200            }
201
202            // bottom line
203            for y in inner[3]..outer[3] {
204                for x in outer[0] + 1..outer[2] {
205                    let index = (y * dst_rs) + (x * dst_c);
206                    dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
207                }
208            }
209        }
210        Ok(())
211    }
212
213    /// Materialize segmentation masks from proto data into `Vec<Segmentation>`.
214    ///
215    /// This is the CPU-side decode step of the hybrid mask rendering path:
216    /// call this to get pre-decoded masks, then pass them to
217    /// [`draw_decoded_masks`](crate::ImageProcessorTrait::draw_decoded_masks) for GPU overlay.
218    /// Benchmarks show this hybrid path (CPU decode + GL overlay) is faster
219    /// than the fused GPU `draw_proto_masks` on all tested platforms.
220    ///
221    /// Optimized: fused dequantization + dot product avoids a 3.1MB f32
222    /// allocation for the full proto tensor. Uses fast sigmoid approximation.
223    pub fn materialize_segmentations(
224        &self,
225        detect: &[crate::DetectBox],
226        proto_data: &crate::ProtoData,
227        letterbox: Option<[f32; 4]>,
228    ) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
229        use edgefirst_tensor::{DType, TensorMapTrait, TensorTrait};
230
231        let _span = tracing::trace_span!(
232            "image.materialize_masks",
233            mode = "proto",
234            n_detections = detect.len(),
235        )
236        .entered();
237
238        if detect.is_empty() {
239            return Ok(Vec::new());
240        }
241        let proto_shape = proto_data.protos.shape();
242        if proto_shape.len() != 3 {
243            return Err(crate::Error::InvalidShape(format!(
244                "protos tensor must be rank-3, got {proto_shape:?}"
245            )));
246        }
247        // Interpret shape based on physical layout.
248        let (proto_h, proto_w, num_protos) = match proto_data.layout {
249            edgefirst_decoder::ProtoLayout::Nhwc => {
250                (proto_shape[0], proto_shape[1], proto_shape[2])
251            }
252            edgefirst_decoder::ProtoLayout::Nchw => {
253                (proto_shape[1], proto_shape[2], proto_shape[0])
254            }
255        };
256        let coeff_shape = proto_data.mask_coefficients.shape();
257        if coeff_shape.len() != 2 || coeff_shape[1] != num_protos {
258            return Err(crate::Error::InvalidShape(format!(
259                "mask_coefficients shape {coeff_shape:?} incompatible with protos \
260                 {proto_shape:?} (expected [N, {num_protos}])"
261            )));
262        }
263        if coeff_shape[0] == 0 {
264            return Ok(Vec::new());
265        }
266        if coeff_shape[0] != detect.len() {
267            return Err(crate::Error::Internal(format!(
268                "mask_coefficients rows {} != detection count {}",
269                coeff_shape[0],
270                detect.len()
271            )));
272        }
273
274        // Precompute inverse letterbox scale for output-coord conversion.
275        let (lx0, inv_lw, ly0, inv_lh) = match letterbox {
276            Some([lx0, ly0, lx1, ly1]) => {
277                let lw = lx1 - lx0;
278                let lh = ly1 - ly0;
279                (
280                    lx0,
281                    if lw > 0.0 { 1.0 / lw } else { 1.0 },
282                    ly0,
283                    if lh > 0.0 { 1.0 / lh } else { 1.0 },
284                )
285            }
286            None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
287        };
288
289        // Fast integer path: when both coefficients and protos are I8 with
290        // per-tensor quantization, use the all-integer kernel (same dot
291        // product infrastructure as the scaled path, but at proto resolution
292        // without bilinear upsampling). Output is binary {0, 255}.
293        // Falls through to the general f32 dequant path for per-channel
294        // quantization or other unsupported modes.
295        if proto_data.mask_coefficients.dtype() == DType::I8
296            && proto_data.protos.dtype() == DType::I8
297        {
298            let coeff_t = proto_data
299                .mask_coefficients
300                .as_i8()
301                .expect("I8 coefficients");
302            let coeff_m = coeff_t.map()?;
303            let coeff_quant = coeff_t.quantization().ok_or_else(|| {
304                crate::Error::InvalidShape(
305                    "I8 mask_coefficients require quantization metadata".into(),
306                )
307            })?;
308            let proto_t = proto_data.protos.as_i8().expect("I8 protos");
309            let proto_m = proto_t.map()?;
310            let proto_quant = proto_t.quantization().ok_or_else(|| {
311                crate::Error::InvalidShape("I8 protos require quantization metadata".into())
312            })?;
313            match proto_segmentations_i8_i8(
314                detect,
315                coeff_m.as_slice(),
316                coeff_quant,
317                proto_m.as_slice(),
318                proto_quant,
319                proto_h,
320                proto_w,
321                num_protos,
322                lx0,
323                inv_lw,
324                ly0,
325                inv_lh,
326                proto_data.layout,
327            ) {
328                Ok(result) => return Ok(result),
329                Err(crate::Error::NotSupported(_)) => {
330                    // Fall through to the general f32 dequant path below for
331                    // per-channel quantization and other unsupported modes.
332                }
333                Err(e) => return Err(e),
334            }
335        }
336
337        // Fast i16×i8 integer path: i16 coefficients × i8 protos → i32 dot.
338        if proto_data.mask_coefficients.dtype() == DType::I16
339            && proto_data.protos.dtype() == DType::I8
340        {
341            let coeff_t = proto_data
342                .mask_coefficients
343                .as_i16()
344                .expect("I16 coefficients");
345            let coeff_m = coeff_t.map()?;
346            // Skip the integer fast path when coefficient quantization is
347            // absent — the f32 fallback below handles raw i16 by widening.
348            if let Some(coeff_quant) = coeff_t.quantization() {
349                let proto_t = proto_data.protos.as_i8().expect("I8 protos");
350                let proto_m = proto_t.map()?;
351                let proto_quant = proto_t.quantization().ok_or_else(|| {
352                    crate::Error::InvalidShape("I8 protos require quantization metadata".into())
353                })?;
354                match proto_segmentations_i16_i8(
355                    detect,
356                    coeff_m.as_slice(),
357                    coeff_quant,
358                    proto_m.as_slice(),
359                    proto_quant,
360                    proto_h,
361                    proto_w,
362                    num_protos,
363                    lx0,
364                    inv_lw,
365                    ly0,
366                    inv_lh,
367                    proto_data.layout,
368                ) {
369                    Ok(result) => return Ok(result),
370                    Err(crate::Error::NotSupported(_)) => {
371                        // Fall through to the general f32 dequant path.
372                    }
373                    Err(e) => return Err(e),
374                }
375            }
376        }
377
378        // Coefficients may be F32 (from f32 models), F16 (from fp16 models),
379        // I8 (from quantized models — kept raw with quantization), or I16.
380        // For the mask kernel we always need an f32 view (the multiply-accumulate
381        // is done in f32 for precision). Map once and widen once outside the loop.
382        if proto_data.layout == edgefirst_decoder::ProtoLayout::Nchw
383            && proto_data.protos.dtype() != DType::I8
384        {
385            return Err(crate::Error::NotSupported(
386                "NCHW proto layout with non-I8 protos is not supported in the f32 fallback path"
387                    .into(),
388            ));
389        }
390        let coeff_f32_storage: Vec<f32>;
391        let coeff_f32_slice: &[f32] = match proto_data.mask_coefficients.dtype() {
392            DType::F32 => {
393                let t = proto_data
394                    .mask_coefficients
395                    .as_f32()
396                    .expect("dtype matched F32");
397                let m = t.map()?;
398                coeff_f32_storage = m.as_slice().to_vec();
399                &coeff_f32_storage[..]
400            }
401            DType::F16 => {
402                let t = proto_data
403                    .mask_coefficients
404                    .as_f16()
405                    .expect("dtype matched F16");
406                let m = t.map()?;
407                coeff_f32_storage = m.as_slice().iter().map(|v| v.to_f32()).collect();
408                &coeff_f32_storage[..]
409            }
410            DType::I8 => {
411                let t = proto_data
412                    .mask_coefficients
413                    .as_i8()
414                    .expect("dtype matched I8");
415                let m = t.map()?;
416                coeff_f32_storage = if let Some(q) = t.quantization() {
417                    use edgefirst_tensor::QuantMode;
418                    let (scale, zp) = match q.mode() {
419                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
420                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
421                        other => {
422                            return Err(crate::Error::NotSupported(format!(
423                                "I8 mask_coefficients quantization mode {other:?} not supported"
424                            )));
425                        }
426                    };
427                    m.as_slice()
428                        .iter()
429                        .map(|&v| (v as f32 - zp) * scale)
430                        .collect()
431                } else {
432                    m.as_slice().iter().map(|&v| v as f32).collect()
433                };
434                &coeff_f32_storage[..]
435            }
436            DType::I16 => {
437                let t = proto_data
438                    .mask_coefficients
439                    .as_i16()
440                    .expect("dtype matched I16");
441                let m = t.map()?;
442                coeff_f32_storage = if let Some(q) = t.quantization() {
443                    use edgefirst_tensor::QuantMode;
444                    let (scale, zp) = match q.mode() {
445                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
446                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
447                        other => {
448                            return Err(crate::Error::NotSupported(format!(
449                                "I16 mask_coefficients quantization mode {other:?} not supported"
450                            )));
451                        }
452                    };
453                    m.as_slice()
454                        .iter()
455                        .map(|&v| (v as f32 - zp) * scale)
456                        .collect()
457                } else {
458                    m.as_slice().iter().map(|&v| v as f32).collect()
459                };
460                &coeff_f32_storage[..]
461            }
462            other => {
463                return Err(crate::Error::InvalidShape(format!(
464                    "mask_coefficients dtype {other:?} not supported; expected F32, F16, I8, or I16"
465                )));
466            }
467        };
468
469        // Hoist the proto tensor map() out of the per-detection loop so the
470        // map-guard is acquired once. Then dispatch per-dtype via a helper
471        // that runs the per-detection kernels in parallel across detections
472        // via rayon. This restores the parallelism that PR #54 added and
473        // PR #51 (EDGEAI-1244 f16 refactor) inadvertently removed.
474        match proto_data.protos.dtype() {
475            DType::I8 => {
476                let t = proto_data.protos.as_i8().expect("dtype matched I8");
477                let quant = t.quantization().ok_or_else(|| {
478                    crate::Error::InvalidShape("I8 protos require quantization metadata".into())
479                })?;
480                let m = t.map()?;
481                let src_slice = m.as_slice();
482                let transposed_storage =
483                    if proto_data.layout == edgefirst_decoder::ProtoLayout::Nchw {
484                        // Guard the proto geometry before the per-plane slice so
485                        // an undersized or overflowing (e.g. imported `from_fd`)
486                        // proto tensor yields `InvalidShape` instead of a wrapped
487                        // `hw` slipping past the bounds check (panic/SIGBUS).
488                        let hw = proto_h.checked_mul(proto_w).ok_or_else(|| {
489                            crate::Error::InvalidShape(format!(
490                                "proto plane size overflow (proto_h={proto_h}, proto_w={proto_w})"
491                            ))
492                        })?;
493                        let need = hw.checked_mul(num_protos).ok_or_else(|| {
494                            crate::Error::InvalidShape(format!(
495                                "proto NCHW size overflow (hw={hw}, n={num_protos})"
496                            ))
497                        })?;
498                        if src_slice.len() < need {
499                            return Err(crate::Error::InvalidShape(format!(
500                                "proto buffer {} bytes < {need} (proto_h={proto_h}, \
501                                 proto_w={proto_w}, num_protos={num_protos})",
502                                src_slice.len()
503                            )));
504                        }
505                        let mut nhwc = vec![0i8; need];
506                        for c in 0..num_protos {
507                            let plane = &src_slice[c * hw..(c + 1) * hw];
508                            for px in 0..hw {
509                                nhwc[px * num_protos + c] = plane[px];
510                            }
511                        }
512                        Some(nhwc)
513                    } else {
514                        None
515                    };
516                let protos_slice = transposed_storage.as_deref().unwrap_or(src_slice);
517                detect
518                    .par_iter()
519                    .enumerate()
520                    .map(|(i, det)| {
521                        let coeff = &coeff_f32_slice[i * num_protos..(i + 1) * num_protos];
522                        let (x0, y0, x1, y1, roi_w, roi_h) =
523                            bbox_to_proto_roi(det, proto_w, proto_h);
524                        let mask = fused_dequant_dot_sign_i8_slice(
525                            protos_slice,
526                            coeff,
527                            quant,
528                            proto_h,
529                            proto_w,
530                            y0,
531                            x0,
532                            roi_h,
533                            roi_w,
534                            num_protos,
535                        )?;
536                        Ok(seg_from_roi(
537                            mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
538                        ))
539                    })
540                    .collect()
541            }
542            DType::F32 => {
543                let t = proto_data.protos.as_f32().expect("dtype matched F32");
544                let m = t.map()?;
545                let protos_slice = m.as_slice();
546                detect
547                    .par_iter()
548                    .enumerate()
549                    .map(|(i, det)| {
550                        let coeff = &coeff_f32_slice[i * num_protos..(i + 1) * num_protos];
551                        let (x0, y0, x1, y1, roi_w, roi_h) =
552                            bbox_to_proto_roi(det, proto_w, proto_h);
553                        let mask = fused_dot_sign_f32_slice(
554                            protos_slice,
555                            coeff,
556                            proto_h,
557                            proto_w,
558                            y0,
559                            x0,
560                            roi_h,
561                            roi_w,
562                            num_protos,
563                        );
564                        Ok(seg_from_roi(
565                            mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
566                        ))
567                    })
568                    .collect()
569            }
570            DType::F16 => {
571                let t = proto_data.protos.as_f16().expect("dtype matched F16");
572                let m = t.map()?;
573                let protos_slice = m.as_slice();
574                detect
575                    .par_iter()
576                    .enumerate()
577                    .map(|(i, det)| {
578                        let coeff = &coeff_f32_slice[i * num_protos..(i + 1) * num_protos];
579                        let (x0, y0, x1, y1, roi_w, roi_h) =
580                            bbox_to_proto_roi(det, proto_w, proto_h);
581                        let mask = fused_dot_sign_f16_slice(
582                            protos_slice,
583                            coeff,
584                            proto_h,
585                            proto_w,
586                            y0,
587                            x0,
588                            roi_h,
589                            roi_w,
590                            num_protos,
591                        );
592                        Ok(seg_from_roi(
593                            mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
594                        ))
595                    })
596                    .collect()
597            }
598            other => Err(crate::Error::InvalidShape(format!(
599                "proto tensor dtype {other:?} not supported"
600            ))),
601        }
602    }
603
604    /// Produce per-detection masks at `(width, height)` pixel resolution by
605    /// upsampling the full proto plane once then cropping per bbox. Each
606    /// `det.bbox` is assumed to be in model-input normalized coordinates
607    /// (the convention used by the decoder output); when `letterbox` is
608    /// `Some`, `(width, height)` are original-content pixel dims and the
609    /// inverse letterbox transform is applied to both the bbox (for the
610    /// crop region and returned `Segmentation` metadata) and each output
611    /// pixel (for proto-plane sampling). Mask values are binary
612    /// `uint8 {0, 255}` after thresholding sigmoid > 0.5.
613    ///
614    /// Used by [`ImageProcessor::materialize_masks`] when the caller selects
615    /// [`MaskResolution::Scaled`](crate::MaskResolution::Scaled).
616    pub fn materialize_scaled_segmentations(
617        &self,
618        detect: &[crate::DetectBox],
619        proto_data: &crate::ProtoData,
620        letterbox: Option<[f32; 4]>,
621        width: u32,
622        height: u32,
623    ) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
624        use edgefirst_tensor::{DType, TensorMapTrait, TensorTrait};
625
626        let _span = tracing::trace_span!(
627            "image.materialize_masks",
628            mode = "scaled",
629            n_detections = detect.len(),
630            width,
631            height,
632        )
633        .entered();
634
635        if detect.is_empty() {
636            return Ok(Vec::new());
637        }
638        if width == 0 || height == 0 {
639            return Err(crate::Error::InvalidShape(
640                "Scaled mask width/height must be positive".into(),
641            ));
642        }
643        let proto_shape = proto_data.protos.shape();
644        if proto_shape.len() != 3 {
645            return Err(crate::Error::InvalidShape(format!(
646                "protos tensor must be rank-3, got {proto_shape:?}"
647            )));
648        }
649        // Interpret shape based on physical layout.
650        let (proto_h, proto_w, num_protos) = match proto_data.layout {
651            edgefirst_decoder::ProtoLayout::Nhwc => {
652                (proto_shape[0], proto_shape[1], proto_shape[2])
653            }
654            edgefirst_decoder::ProtoLayout::Nchw => {
655                (proto_shape[1], proto_shape[2], proto_shape[0])
656            }
657        };
658        let coeff_shape = proto_data.mask_coefficients.shape();
659        if coeff_shape.len() != 2 || coeff_shape[1] != num_protos {
660            return Err(crate::Error::InvalidShape(format!(
661                "mask_coefficients shape {coeff_shape:?} incompatible with protos \
662                 {proto_shape:?}"
663            )));
664        }
665        if coeff_shape[0] == 0 {
666            return Ok(Vec::new());
667        }
668        if coeff_shape[0] != detect.len() {
669            return Err(crate::Error::Internal(format!(
670                "mask_coefficients rows {} != detection count {}",
671                coeff_shape[0],
672                detect.len()
673            )));
674        }
675
676        // Fast integer path: when both coefficients and protos are I8, use
677        // the all-integer kernel (i8×i8→i32 dot product, sign-shortcut
678        // bilinear). No floating-point conversion at all.
679        // Falls through to the general f32 dequant path for per-channel
680        // quantization or other unsupported modes.
681        if proto_data.mask_coefficients.dtype() == DType::I8
682            && proto_data.protos.dtype() == DType::I8
683        {
684            let coeff_t = proto_data
685                .mask_coefficients
686                .as_i8()
687                .expect("I8 coefficients");
688            let coeff_m = coeff_t.map()?;
689            let coeff_quant = coeff_t.quantization().ok_or_else(|| {
690                crate::Error::InvalidShape(
691                    "I8 mask_coefficients require quantization metadata".into(),
692                )
693            })?;
694            let proto_t = proto_data.protos.as_i8().expect("I8 protos");
695            let proto_m = proto_t.map()?;
696            let proto_quant = proto_t.quantization().ok_or_else(|| {
697                crate::Error::InvalidShape("I8 protos require quantization metadata".into())
698            })?;
699            match scaled_segmentations_i8_i8(
700                detect,
701                coeff_m.as_slice(),
702                coeff_quant,
703                proto_m.as_slice(),
704                proto_quant,
705                proto_h,
706                proto_w,
707                num_protos,
708                letterbox,
709                width,
710                height,
711                proto_data.layout,
712            ) {
713                Ok(result) => return Ok(result),
714                Err(crate::Error::NotSupported(_)) => {
715                    // Fall through to the general f32 dequant path below for
716                    // per-channel quantization and other unsupported modes.
717                }
718                Err(e) => return Err(e),
719            }
720        }
721
722        // Fast i16×i8 integer path: i16 coefficients × i8 protos.
723        if proto_data.mask_coefficients.dtype() == DType::I16
724            && proto_data.protos.dtype() == DType::I8
725        {
726            let coeff_t = proto_data
727                .mask_coefficients
728                .as_i16()
729                .expect("I16 coefficients");
730            let coeff_m = coeff_t.map()?;
731            // Skip the integer fast path when coefficient quantization is
732            // absent — the f32 fallback below handles raw i16 by widening.
733            if let Some(coeff_quant) = coeff_t.quantization() {
734                let proto_t = proto_data.protos.as_i8().expect("I8 protos");
735                let proto_m = proto_t.map()?;
736                let proto_quant = proto_t.quantization().ok_or_else(|| {
737                    crate::Error::InvalidShape("I8 protos require quantization metadata".into())
738                })?;
739                match scaled_segmentations_i16_i8(
740                    detect,
741                    coeff_m.as_slice(),
742                    coeff_quant,
743                    proto_m.as_slice(),
744                    proto_quant,
745                    proto_h,
746                    proto_w,
747                    num_protos,
748                    letterbox,
749                    width,
750                    height,
751                    proto_data.layout,
752                ) {
753                    Ok(result) => return Ok(result),
754                    Err(crate::Error::NotSupported(_)) => {}
755                    Err(e) => return Err(e),
756                }
757            }
758        }
759
760        // Fallback: widen coefficients to f32 for the float-path kernels.
761        if proto_data.layout == edgefirst_decoder::ProtoLayout::Nchw
762            && proto_data.protos.dtype() != DType::I8
763        {
764            return Err(crate::Error::NotSupported(
765                "NCHW proto layout with non-I8 protos is not supported in the f32 fallback path"
766                    .into(),
767            ));
768        }
769        let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
770            DType::F32 => {
771                let t = proto_data.mask_coefficients.as_f32().expect("F32");
772                let m = t.map()?;
773                m.as_slice().to_vec()
774            }
775            DType::F16 => {
776                let t = proto_data.mask_coefficients.as_f16().expect("F16");
777                let m = t.map()?;
778                m.as_slice().iter().map(|v| v.to_f32()).collect()
779            }
780            DType::I8 => {
781                // Dequantize I8 coefficients to f32 for the float proto path.
782                let t = proto_data.mask_coefficients.as_i8().expect("I8");
783                let m = t.map()?;
784                let q = t.quantization().ok_or_else(|| {
785                    crate::Error::InvalidShape(
786                        "I8 mask_coefficients require quantization metadata".into(),
787                    )
788                })?;
789                use edgefirst_tensor::QuantMode;
790                let (scale, zp) = match q.mode() {
791                    QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
792                    QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
793                    _ => {
794                        return Err(crate::Error::NotSupported(
795                            "per-channel mask_coefficients not supported".into(),
796                        ))
797                    }
798                };
799                m.as_slice()
800                    .iter()
801                    .map(|&v| (v as f32 - zp) * scale)
802                    .collect()
803            }
804            DType::I16 => {
805                let t = proto_data.mask_coefficients.as_i16().expect("I16");
806                let m = t.map()?;
807                if let Some(q) = t.quantization() {
808                    use edgefirst_tensor::QuantMode;
809                    let (scale, zp) = match q.mode() {
810                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
811                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
812                        other => {
813                            return Err(crate::Error::NotSupported(format!(
814                                "I16 mask_coefficients quantization mode {other:?} not supported"
815                            )))
816                        }
817                    };
818                    m.as_slice()
819                        .iter()
820                        .map(|&v| (v as f32 - zp) * scale)
821                        .collect()
822                } else {
823                    m.as_slice().iter().map(|&v| v as f32).collect()
824                }
825            }
826            other => {
827                return Err(crate::Error::InvalidShape(format!(
828                    "mask_coefficients dtype {other:?} not supported"
829                )));
830            }
831        };
832
833        match proto_data.protos.dtype() {
834            DType::F32 => {
835                let t = proto_data.protos.as_f32().expect("F32");
836                let m = t.map()?;
837                scaled_segmentations_f32_slice(
838                    detect,
839                    &coeff_f32,
840                    m.as_slice(),
841                    proto_h,
842                    proto_w,
843                    num_protos,
844                    letterbox,
845                    width,
846                    height,
847                )
848            }
849            DType::F16 => {
850                let t = proto_data.protos.as_f16().expect("F16");
851                let m = t.map()?;
852                scaled_segmentations_f16_slice(
853                    detect,
854                    &coeff_f32,
855                    m.as_slice(),
856                    proto_h,
857                    proto_w,
858                    num_protos,
859                    letterbox,
860                    width,
861                    height,
862                )
863            }
864            DType::I8 => {
865                let t = proto_data.protos.as_i8().expect("I8");
866                let m = t.map()?;
867                let quant = t.quantization().ok_or_else(|| {
868                    crate::Error::InvalidShape("I8 protos require quantization metadata".into())
869                })?;
870                let src_slice = m.as_slice();
871                let transposed_storage =
872                    if proto_data.layout == edgefirst_decoder::ProtoLayout::Nchw {
873                        // Guard the proto geometry before the per-plane slice so
874                        // an undersized or overflowing (e.g. imported `from_fd`)
875                        // proto tensor yields `InvalidShape` instead of a wrapped
876                        // `hw` slipping past the bounds check (panic/SIGBUS).
877                        let hw = proto_h.checked_mul(proto_w).ok_or_else(|| {
878                            crate::Error::InvalidShape(format!(
879                                "proto plane size overflow (proto_h={proto_h}, proto_w={proto_w})"
880                            ))
881                        })?;
882                        let need = hw.checked_mul(num_protos).ok_or_else(|| {
883                            crate::Error::InvalidShape(format!(
884                                "proto NCHW size overflow (hw={hw}, n={num_protos})"
885                            ))
886                        })?;
887                        if src_slice.len() < need {
888                            return Err(crate::Error::InvalidShape(format!(
889                                "proto buffer {} bytes < {need} (proto_h={proto_h}, \
890                                 proto_w={proto_w}, num_protos={num_protos})",
891                                src_slice.len()
892                            )));
893                        }
894                        let mut nhwc = vec![0i8; need];
895                        for c in 0..num_protos {
896                            let plane = &src_slice[c * hw..(c + 1) * hw];
897                            for px in 0..hw {
898                                nhwc[px * num_protos + c] = plane[px];
899                            }
900                        }
901                        Some(nhwc)
902                    } else {
903                        None
904                    };
905                let protos_slice = transposed_storage.as_deref().unwrap_or(src_slice);
906                scaled_segmentations_i8_slice(
907                    detect,
908                    &coeff_f32,
909                    protos_slice,
910                    proto_h,
911                    proto_w,
912                    num_protos,
913                    quant,
914                    letterbox,
915                    width,
916                    height,
917                )
918            }
919            other => Err(crate::Error::InvalidShape(format!(
920                "proto tensor dtype {other:?} not supported"
921            ))),
922        }
923    }
924}
925
926// =============================================================================
927// Slice-native fused kernels.
928//
929// All kernels take row-major `[H, W, num_protos]` proto slices + `&[f32]`
930// coefficients (widened once from the source dtype at the materialize entry
931// point). Per-dtype variants exist for i8 (with on-the-fly dequant using a
932// tensor-level `Quantization`), f32, and f16; f16 widens to f32 per-element
933// at the FMA site via `half::f16::to_f32()`.
934//
935// On ARMv8.2-FP16 this compiles to `fcvt`; on Cortex-A53 and non-F16C x86 it
936// becomes a soft-float helper.
937// =============================================================================
938
939/// Map a detection bbox in normalised letterboxed coords to its ROI in
940/// the proto plane (floor xmin/ymin, ceil xmax/ymax, clamp to plane bounds).
941/// Returns `(x0, y0, x1, y1, roi_w, roi_h)` where roi_w/h are guaranteed ≥ 1.
942fn bbox_to_proto_roi(
943    det: &DetectBox,
944    proto_w: usize,
945    proto_h: usize,
946) -> (usize, usize, usize, usize, usize, usize) {
947    let bbox = det.bbox.to_canonical();
948    let xmin = bbox.xmin.clamp(0.0, 1.0);
949    let ymin = bbox.ymin.clamp(0.0, 1.0);
950    let xmax = bbox.xmax.clamp(0.0, 1.0);
951    let ymax = bbox.ymax.clamp(0.0, 1.0);
952    let x0 = ((xmin * proto_w as f32) as usize).min(proto_w.saturating_sub(1));
953    let y0 = ((ymin * proto_h as f32) as usize).min(proto_h.saturating_sub(1));
954    let x1 = ((xmax * proto_w as f32).ceil() as usize).min(proto_w);
955    let y1 = ((ymax * proto_h as f32).ceil() as usize).min(proto_h);
956    let roi_w = x1.saturating_sub(x0).max(1);
957    let roi_h = y1.saturating_sub(y0).max(1);
958    (x0, y0, x1, y1, roi_w, roi_h)
959}
960
961/// Build a `Segmentation` from a per-detection mask + the ROI bounds in
962/// proto coords. Applies the inverse letterbox transform to express the
963/// segmentation bbox in original-image-content normalised space.
964#[allow(clippy::too_many_arguments)]
965fn seg_from_roi(
966    mask: ndarray::Array3<u8>,
967    x0: usize,
968    y0: usize,
969    x1: usize,
970    y1: usize,
971    proto_w: usize,
972    proto_h: usize,
973    lx0: f32,
974    inv_lw: f32,
975    ly0: f32,
976    inv_lh: f32,
977) -> edgefirst_decoder::Segmentation {
978    let seg_xmin = ((x0 as f32 / proto_w as f32) - lx0) * inv_lw;
979    let seg_ymin = ((y0 as f32 / proto_h as f32) - ly0) * inv_lh;
980    let seg_xmax = ((x1 as f32 / proto_w as f32) - lx0) * inv_lw;
981    let seg_ymax = ((y1 as f32 / proto_h as f32) - ly0) * inv_lh;
982    edgefirst_decoder::Segmentation {
983        xmin: seg_xmin.clamp(0.0, 1.0),
984        ymin: seg_ymin.clamp(0.0, 1.0),
985        xmax: seg_xmax.clamp(0.0, 1.0),
986        ymax: seg_ymax.clamp(0.0, 1.0),
987        segmentation: mask,
988    }
989}
990
991// =============================================================================
992// Integer-domain proto-resolution kernel: i8 coefficients × i8 protos → i32
993// → sign threshold → binary {0, 255}.
994//
995// Reuses the same dot product infrastructure as the scaled path (NEON sdot on
996// A55+, smull+sadalp on A53, scalar fallback on x86). Since proto-resolution
997// produces masks at the native proto grid (~30×30 per ROI), there is no
998// bilinear upsampling — just a direct sign threshold per pixel.
999// =============================================================================
1000
1001/// Proto-resolution mask materialization using integer-domain math.
1002///
1003/// For each detection, computes the i8×i8 dot product at every proto-ROI pixel,
1004/// applies the zero-point correction, and thresholds at sign(logit) → {0, 255}.
1005/// Supports both NHWC and NCHW proto layouts.
1006#[allow(clippy::too_many_arguments)]
1007fn proto_segmentations_i8_i8(
1008    detect: &[crate::DetectBox],
1009    coeff_all: &[i8],
1010    coeff_quant: &edgefirst_tensor::Quantization,
1011    protos: &[i8],
1012    proto_quant: &edgefirst_tensor::Quantization,
1013    proto_h: usize,
1014    proto_w: usize,
1015    num_protos: usize,
1016    lx0: f32,
1017    inv_lw: f32,
1018    ly0: f32,
1019    inv_lh: f32,
1020    layout: edgefirst_decoder::ProtoLayout,
1021) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
1022    use edgefirst_tensor::QuantMode;
1023
1024    let _span = tracing::trace_span!(
1025        "image.materialize_masks.kernel_i8",
1026        n = detect.len(),
1027        proto_h,
1028        proto_w,
1029        num_protos,
1030        ?layout,
1031    )
1032    .entered();
1033
1034    let zp_c: i32 = match coeff_quant.mode() {
1035        QuantMode::PerTensor { zero_point, .. } => zero_point,
1036        QuantMode::PerTensorSymmetric { .. } => 0,
1037        _ => {
1038            return Err(crate::Error::NotSupported(
1039                "per-channel coeff quantization not supported on proto-res i8 path".into(),
1040            ))
1041        }
1042    };
1043    let zp_p: i32 = match proto_quant.mode() {
1044        QuantMode::PerTensor { zero_point, .. } => zero_point,
1045        QuantMode::PerTensorSymmetric { .. } => 0,
1046        _ => {
1047            return Err(crate::Error::NotSupported(
1048                "per-channel proto quantization not supported on proto-res i8 path".into(),
1049            ))
1050        }
1051    };
1052
1053    let hw = proto_h * proto_w;
1054
1055    // Precompute per-pixel proto sums for zero-point correction.
1056    let proto_sums: Vec<i32> = if zp_c != 0 {
1057        match layout {
1058            edgefirst_decoder::ProtoLayout::Nhwc => (0..hw)
1059                .map(|px_idx| {
1060                    let base = px_idx * num_protos;
1061                    protos[base..base + num_protos]
1062                        .iter()
1063                        .map(|&v| v as i32)
1064                        .sum()
1065                })
1066                .collect(),
1067            edgefirst_decoder::ProtoLayout::Nchw => {
1068                let mut sums = vec![0i32; hw];
1069                for c in 0..num_protos {
1070                    let plane = &protos[c * hw..];
1071                    for (px, s) in sums.iter_mut().enumerate() {
1072                        *s += plane[px] as i32;
1073                    }
1074                }
1075                sums
1076            }
1077        }
1078    } else {
1079        Vec::new()
1080    };
1081
1082    #[cfg(target_arch = "aarch64")]
1083    let use_dotprod = std::arch::is_aarch64_feature_detected!("dotprod");
1084
1085    detect
1086        .par_iter()
1087        .enumerate()
1088        .map(|(i, det)| {
1089            let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
1090            let (x0, y0, x1, y1, roi_w, roi_h) = bbox_to_proto_roi(det, proto_w, proto_h);
1091
1092            // Per-detection bias: zp_p·Σc_raw - N·zp_c·zp_p
1093            let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
1094            let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
1095
1096            let mut mask_buf = vec![0u8; roi_h * roi_w];
1097
1098            match layout {
1099                edgefirst_decoder::ProtoLayout::Nhwc => {
1100                    let stride_y = proto_w * num_protos;
1101                    #[cfg(target_arch = "aarch64")]
1102                    {
1103                        if use_dotprod {
1104                            for ly in 0..roi_h {
1105                                let py = y0 + ly;
1106                                let row_base = py * stride_y + x0 * num_protos;
1107                                for lx in 0..roi_w {
1108                                    let pix_base = row_base + lx * num_protos;
1109                                    let proto_px = &protos[pix_base..pix_base + num_protos];
1110                                    let raw_dot = unsafe {
1111                                        dot_i8_neon_dotprod(
1112                                            coeff.as_ptr(),
1113                                            proto_px.as_ptr(),
1114                                            num_protos,
1115                                        )
1116                                    };
1117                                    let correction = if zp_c != 0 {
1118                                        zp_c * proto_sums[py * proto_w + x0 + lx]
1119                                    } else {
1120                                        0
1121                                    };
1122                                    let logit = raw_dot - correction - bias;
1123                                    if logit > 0 {
1124                                        mask_buf[ly * roi_w + lx] = 255;
1125                                    }
1126                                }
1127                            }
1128                        } else {
1129                            for ly in 0..roi_h {
1130                                let py = y0 + ly;
1131                                let row_base = py * stride_y + x0 * num_protos;
1132                                for lx in 0..roi_w {
1133                                    let pix_base = row_base + lx * num_protos;
1134                                    let proto_px = &protos[pix_base..pix_base + num_protos];
1135                                    let raw_dot = unsafe {
1136                                        dot_i8_neon_base(
1137                                            coeff.as_ptr(),
1138                                            proto_px.as_ptr(),
1139                                            num_protos,
1140                                        )
1141                                    };
1142                                    let correction = if zp_c != 0 {
1143                                        zp_c * proto_sums[py * proto_w + x0 + lx]
1144                                    } else {
1145                                        0
1146                                    };
1147                                    let logit = raw_dot - correction - bias;
1148                                    if logit > 0 {
1149                                        mask_buf[ly * roi_w + lx] = 255;
1150                                    }
1151                                }
1152                            }
1153                        }
1154                    }
1155                    #[cfg(not(target_arch = "aarch64"))]
1156                    {
1157                        for ly in 0..roi_h {
1158                            let py = y0 + ly;
1159                            let row_base = py * stride_y + x0 * num_protos;
1160                            for lx in 0..roi_w {
1161                                let pix_base = row_base + lx * num_protos;
1162                                let proto_px = &protos[pix_base..pix_base + num_protos];
1163                                let raw_dot = dot_i8_scalar(coeff, proto_px, num_protos);
1164                                let correction = if zp_c != 0 {
1165                                    zp_c * proto_sums[py * proto_w + x0 + lx]
1166                                } else {
1167                                    0
1168                                };
1169                                let logit = raw_dot - correction - bias;
1170                                if logit > 0 {
1171                                    mask_buf[ly * roi_w + lx] = 255;
1172                                }
1173                            }
1174                        }
1175                    }
1176                }
1177                edgefirst_decoder::ProtoLayout::Nchw => {
1178                    // Channel-major accumulation: for each channel, accumulate
1179                    // coeff[c] * proto[c, py, px] across the ROI. Each channel
1180                    // plane is contiguous, giving excellent sequential read access.
1181                    let mut accum = vec![0i32; roi_h * roi_w];
1182                    for c in 0..num_protos {
1183                        let plane = &protos[c * hw..];
1184                        let coeff_c = coeff[c] as i32;
1185                        for ly in 0..roi_h {
1186                            let py = y0 + ly;
1187                            let row_start = py * proto_w + x0;
1188                            let out_row_start = ly * roi_w;
1189                            for lx in 0..roi_w {
1190                                accum[out_row_start + lx] += coeff_c * plane[row_start + lx] as i32;
1191                            }
1192                        }
1193                    }
1194                    // Apply zero-point correction and threshold.
1195                    for ly in 0..roi_h {
1196                        let py = y0 + ly;
1197                        for lx in 0..roi_w {
1198                            let idx = ly * roi_w + lx;
1199                            let correction = if zp_c != 0 {
1200                                zp_c * proto_sums[py * proto_w + x0 + lx]
1201                            } else {
1202                                0
1203                            };
1204                            let logit = accum[idx] - correction - bias;
1205                            if logit > 0 {
1206                                mask_buf[idx] = 255;
1207                            }
1208                        }
1209                    }
1210                }
1211            }
1212
1213            let mask = ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1214                .expect("mask_buf length matches roi_h * roi_w");
1215            Ok(seg_from_roi(
1216                mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
1217            ))
1218        })
1219        .collect()
1220}
1221
1222// =============================================================================
1223#[allow(clippy::too_many_arguments)]
1224fn proto_segmentations_i16_i8(
1225    detect: &[crate::DetectBox],
1226    coeff_all: &[i16],
1227    coeff_quant: &edgefirst_tensor::Quantization,
1228    protos: &[i8],
1229    proto_quant: &edgefirst_tensor::Quantization,
1230    proto_h: usize,
1231    proto_w: usize,
1232    num_protos: usize,
1233    lx0: f32,
1234    inv_lw: f32,
1235    ly0: f32,
1236    inv_lh: f32,
1237    layout: edgefirst_decoder::ProtoLayout,
1238) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
1239    use edgefirst_tensor::QuantMode;
1240
1241    let _span = tracing::trace_span!(
1242        "image.materialize_masks.kernel_i16xi8",
1243        n = detect.len(),
1244        proto_h,
1245        proto_w,
1246        num_protos,
1247        ?layout,
1248    )
1249    .entered();
1250
1251    let zp_c: i32 = match coeff_quant.mode() {
1252        QuantMode::PerTensor { zero_point, .. } => zero_point,
1253        QuantMode::PerTensorSymmetric { .. } => 0,
1254        _ => {
1255            return Err(crate::Error::NotSupported(
1256                "per-channel coeff quantization not supported on proto-res i16 path".into(),
1257            ))
1258        }
1259    };
1260    let zp_p: i32 = match proto_quant.mode() {
1261        QuantMode::PerTensor { zero_point, .. } => zero_point,
1262        QuantMode::PerTensorSymmetric { .. } => 0,
1263        _ => {
1264            return Err(crate::Error::NotSupported(
1265                "per-channel proto quantization not supported on proto-res i8 path".into(),
1266            ))
1267        }
1268    };
1269
1270    let hw = proto_h * proto_w;
1271
1272    // Precompute per-pixel proto sums for zero-point correction.
1273    let proto_sums: Vec<i32> = if zp_c != 0 {
1274        match layout {
1275            edgefirst_decoder::ProtoLayout::Nhwc => (0..hw)
1276                .map(|px_idx| {
1277                    let base = px_idx * num_protos;
1278                    protos[base..base + num_protos]
1279                        .iter()
1280                        .map(|&v| v as i32)
1281                        .sum()
1282                })
1283                .collect(),
1284            edgefirst_decoder::ProtoLayout::Nchw => {
1285                let mut sums = vec![0i32; hw];
1286                for c in 0..num_protos {
1287                    let plane = &protos[c * hw..];
1288                    for (px, s) in sums.iter_mut().enumerate() {
1289                        *s += plane[px] as i32;
1290                    }
1291                }
1292                sums
1293            }
1294        }
1295    } else {
1296        Vec::new()
1297    };
1298
1299    detect
1300        .par_iter()
1301        .enumerate()
1302        .map(|(i, det)| {
1303            let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
1304            let (x0, y0, x1, y1, roi_w, roi_h) = bbox_to_proto_roi(det, proto_w, proto_h);
1305
1306            // Per-detection bias: zp_p·Σc_raw - N·zp_c·zp_p
1307            let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
1308            let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
1309
1310            let mut mask_buf = vec![0u8; roi_h * roi_w];
1311
1312            match layout {
1313                edgefirst_decoder::ProtoLayout::Nhwc => {
1314                    let stride_y = proto_w * num_protos;
1315                    #[cfg(target_arch = "aarch64")]
1316                    {
1317                        for ly in 0..roi_h {
1318                            let py = y0 + ly;
1319                            let row_base = py * stride_y + x0 * num_protos;
1320                            for lx in 0..roi_w {
1321                                let pix_base = row_base + lx * num_protos;
1322                                let proto_px = &protos[pix_base..pix_base + num_protos];
1323                                let raw_dot = unsafe {
1324                                    dot_i16_i8_neon(coeff.as_ptr(), proto_px.as_ptr(), num_protos)
1325                                };
1326                                let correction = if zp_c != 0 {
1327                                    zp_c * proto_sums[py * proto_w + x0 + lx]
1328                                } else {
1329                                    0
1330                                };
1331                                let logit = raw_dot - correction - bias;
1332                                if logit > 0 {
1333                                    mask_buf[ly * roi_w + lx] = 255;
1334                                }
1335                            }
1336                        }
1337                    }
1338                    #[cfg(not(target_arch = "aarch64"))]
1339                    {
1340                        for ly in 0..roi_h {
1341                            let py = y0 + ly;
1342                            let row_base = py * stride_y + x0 * num_protos;
1343                            for lx in 0..roi_w {
1344                                let pix_base = row_base + lx * num_protos;
1345                                let proto_px = &protos[pix_base..pix_base + num_protos];
1346                                let raw_dot = dot_i16_i8_scalar(coeff, proto_px, num_protos);
1347                                let correction = if zp_c != 0 {
1348                                    zp_c * proto_sums[py * proto_w + x0 + lx]
1349                                } else {
1350                                    0
1351                                };
1352                                let logit = raw_dot - correction - bias;
1353                                if logit > 0 {
1354                                    mask_buf[ly * roi_w + lx] = 255;
1355                                }
1356                            }
1357                        }
1358                    }
1359                }
1360                edgefirst_decoder::ProtoLayout::Nchw => {
1361                    // Channel-major accumulation: for each channel, accumulate
1362                    // coeff[c] * proto[c, py, px] across the ROI. Each channel
1363                    // plane is contiguous, giving excellent sequential read access.
1364                    let mut accum = vec![0i32; roi_h * roi_w];
1365                    for c in 0..num_protos {
1366                        let plane = &protos[c * hw..];
1367                        let coeff_c = coeff[c] as i32;
1368                        for ly in 0..roi_h {
1369                            let py = y0 + ly;
1370                            let row_start = py * proto_w + x0;
1371                            let out_row_start = ly * roi_w;
1372                            for lx in 0..roi_w {
1373                                accum[out_row_start + lx] += coeff_c * plane[row_start + lx] as i32;
1374                            }
1375                        }
1376                    }
1377                    // Apply zero-point correction and threshold.
1378                    for ly in 0..roi_h {
1379                        let py = y0 + ly;
1380                        for lx in 0..roi_w {
1381                            let idx = ly * roi_w + lx;
1382                            let correction = if zp_c != 0 {
1383                                zp_c * proto_sums[py * proto_w + x0 + lx]
1384                            } else {
1385                                0
1386                            };
1387                            let logit = accum[idx] - correction - bias;
1388                            if logit > 0 {
1389                                mask_buf[idx] = 255;
1390                            }
1391                        }
1392                    }
1393                }
1394            }
1395
1396            let mask = ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1397                .expect("mask_buf length matches roi_h * roi_w");
1398            Ok(seg_from_roi(
1399                mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
1400            ))
1401        })
1402        .collect()
1403}
1404
1405// =============================================================================
1406
1407// Sign-threshold proto-resolution kernels (f32/f16/i8 protos with f32 coeffs).
1408//
1409// These replace the sigmoid-computing kernels for the non-i8×i8 fallback paths.
1410// Since downstream always thresholds at > 127, computing sigmoid is wasteful;
1411// sign(dot) > 0 ⟺ sigmoid(dot) > 0.5 gives the same binary result.
1412// =============================================================================
1413
1414/// f32 protos × f32 coefficients → sign threshold → binary {0, 255}.
1415#[allow(clippy::too_many_arguments)]
1416fn fused_dot_sign_f32_slice(
1417    protos: &[f32],
1418    coeff: &[f32],
1419    _proto_h: usize,
1420    proto_w: usize,
1421    y0: usize,
1422    x0: usize,
1423    roi_h: usize,
1424    roi_w: usize,
1425    num_protos: usize,
1426) -> ndarray::Array3<u8> {
1427    let stride_y = proto_w * num_protos;
1428    let mut mask_buf = vec![0u8; roi_h * roi_w];
1429    for y in 0..roi_h {
1430        let row_base = (y0 + y) * stride_y + x0 * num_protos;
1431        let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1432        for (x, out_px) in out_row.iter_mut().enumerate() {
1433            let base = row_base + x * num_protos;
1434            let mut acc = 0.0_f32;
1435            let mut k = 0;
1436            let chunks = num_protos / 4;
1437            for _ in 0..chunks {
1438                acc += coeff[k] * protos[base + k]
1439                    + coeff[k + 1] * protos[base + k + 1]
1440                    + coeff[k + 2] * protos[base + k + 2]
1441                    + coeff[k + 3] * protos[base + k + 3];
1442                k += 4;
1443            }
1444            while k < num_protos {
1445                acc += coeff[k] * protos[base + k];
1446                k += 1;
1447            }
1448            if acc > 0.0 {
1449                *out_px = 255;
1450            }
1451        }
1452    }
1453    ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1454        .expect("mask_buf length matches roi_h * roi_w")
1455}
1456
1457/// f16 protos × f32 coefficients → sign threshold → binary {0, 255}.
1458///
1459/// Two code paths:
1460///
1461/// 1. **x86_64 + F16C + FMA** — explicit intrinsic kernel using
1462///    `_mm256_cvtph_ps` (8-lane f16→f32 widening) + `_mm256_fmadd_ps`.
1463///
1464/// 2. **Scalar fallback** — loop-unrolled by 4 with `half::f16::to_f32()`.
1465#[allow(clippy::too_many_arguments)]
1466fn fused_dot_sign_f16_slice(
1467    protos: &[half::f16],
1468    coeff: &[f32],
1469    _proto_h: usize,
1470    proto_w: usize,
1471    y0: usize,
1472    x0: usize,
1473    roi_h: usize,
1474    roi_w: usize,
1475    num_protos: usize,
1476) -> ndarray::Array3<u8> {
1477    #[cfg(all(
1478        target_arch = "x86_64",
1479        target_feature = "f16c",
1480        target_feature = "fma"
1481    ))]
1482    {
1483        // SAFETY: target-feature gates guarantee F16C + FMA support.
1484        unsafe {
1485            fused_dot_sign_f16_slice_f16c(protos, coeff, proto_w, y0, x0, roi_h, roi_w, num_protos)
1486        }
1487    }
1488    #[cfg(not(all(
1489        target_arch = "x86_64",
1490        target_feature = "f16c",
1491        target_feature = "fma"
1492    )))]
1493    {
1494        fused_dot_sign_f16_slice_scalar(protos, coeff, proto_w, y0, x0, roi_h, roi_w, num_protos)
1495    }
1496}
1497
1498/// Scalar f16 sign-threshold kernel — loop-unrolled by 4.
1499#[allow(clippy::too_many_arguments)]
1500fn fused_dot_sign_f16_slice_scalar(
1501    protos: &[half::f16],
1502    coeff: &[f32],
1503    proto_w: usize,
1504    y0: usize,
1505    x0: usize,
1506    roi_h: usize,
1507    roi_w: usize,
1508    num_protos: usize,
1509) -> ndarray::Array3<u8> {
1510    let stride_y = proto_w * num_protos;
1511    let mut mask_buf = vec![0u8; roi_h * roi_w];
1512    for y in 0..roi_h {
1513        let row_base = (y0 + y) * stride_y + x0 * num_protos;
1514        let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1515        for (x, out_px) in out_row.iter_mut().enumerate() {
1516            let base = row_base + x * num_protos;
1517            let mut acc = 0.0_f32;
1518            let mut k = 0;
1519            let chunks = num_protos / 4;
1520            for _ in 0..chunks {
1521                acc += coeff[k] * protos[base + k].to_f32()
1522                    + coeff[k + 1] * protos[base + k + 1].to_f32()
1523                    + coeff[k + 2] * protos[base + k + 2].to_f32()
1524                    + coeff[k + 3] * protos[base + k + 3].to_f32();
1525                k += 4;
1526            }
1527            while k < num_protos {
1528                acc += coeff[k] * protos[base + k].to_f32();
1529                k += 1;
1530            }
1531            if acc > 0.0 {
1532                *out_px = 255;
1533            }
1534        }
1535    }
1536    ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1537        .expect("mask_buf length matches roi_h * roi_w")
1538}
1539
1540/// x86_64 F16C + FMA intrinsic kernel for f16 sign-threshold.
1541///
1542/// Uses `_mm256_cvtph_ps` for hardware 8-lane f16→f32 widening and
1543/// `_mm256_fmadd_ps` for fused multiply-add. Only the sign of the
1544/// accumulated dot product is checked (no sigmoid needed).
1545///
1546/// # Safety
1547///
1548/// Caller must ensure the target CPU supports F16C + FMA.
1549#[cfg(all(
1550    target_arch = "x86_64",
1551    target_feature = "f16c",
1552    target_feature = "fma"
1553))]
1554#[allow(clippy::too_many_arguments)]
1555#[target_feature(enable = "f16c,fma,avx")]
1556unsafe fn fused_dot_sign_f16_slice_f16c(
1557    protos: &[half::f16],
1558    coeff: &[f32],
1559    proto_w: usize,
1560    y0: usize,
1561    x0: usize,
1562    roi_h: usize,
1563    roi_w: usize,
1564    num_protos: usize,
1565) -> ndarray::Array3<u8> {
1566    use core::arch::x86_64::{
1567        _mm256_castps256_ps128, _mm256_cvtph_ps, _mm256_extractf128_ps, _mm256_fmadd_ps,
1568        _mm256_loadu_ps, _mm256_setzero_ps, _mm_add_ps, _mm_cvtss_f32, _mm_hadd_ps,
1569        _mm_loadu_si128,
1570    };
1571
1572    let stride_y = proto_w * num_protos;
1573    let chunks8 = num_protos / 8;
1574    let mut mask_buf = vec![0u8; roi_h * roi_w];
1575
1576    for y in 0..roi_h {
1577        let row_base = (y0 + y) * stride_y + x0 * num_protos;
1578        let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1579        for (x, out_px) in out_row.iter_mut().enumerate() {
1580            let base = row_base + x * num_protos;
1581            let mut acc_v = _mm256_setzero_ps();
1582            let mut k = 0;
1583            for _ in 0..chunks8 {
1584                let p_ptr = protos
1585                    .as_ptr()
1586                    .add(base + k)
1587                    .cast::<core::arch::x86_64::__m128i>();
1588                let raw = _mm_loadu_si128(p_ptr);
1589                let widened = _mm256_cvtph_ps(raw);
1590                let coeffs_v = _mm256_loadu_ps(coeff.as_ptr().add(k));
1591                acc_v = _mm256_fmadd_ps(coeffs_v, widened, acc_v);
1592                k += 8;
1593            }
1594            // Horizontal reduce 8 → 1.
1595            let lo = _mm256_castps256_ps128(acc_v);
1596            let hi = _mm256_extractf128_ps::<1>(acc_v);
1597            let sum4 = _mm_add_ps(lo, hi);
1598            let sum2 = _mm_hadd_ps(sum4, sum4);
1599            let sum1 = _mm_hadd_ps(sum2, sum2);
1600            let mut acc = _mm_cvtss_f32(sum1);
1601
1602            // Scalar tail for num_protos % 8.
1603            while k < num_protos {
1604                acc += coeff[k] * protos[base + k].to_f32();
1605                k += 1;
1606            }
1607
1608            if acc > 0.0 {
1609                *out_px = 255;
1610            }
1611        }
1612    }
1613    ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1614        .expect("mask_buf length matches roi_h * roi_w")
1615}
1616
1617/// i8 protos (with quant) × f32 coefficients → sign threshold → binary {0, 255}.
1618/// Fallback for per-channel quant or mixed-dtype cases where the i8×i8 fast path
1619/// doesn't apply.
1620#[allow(clippy::too_many_arguments)]
1621fn fused_dequant_dot_sign_i8_slice(
1622    protos: &[i8],
1623    coeff: &[f32],
1624    quant: &edgefirst_tensor::Quantization,
1625    _proto_h: usize,
1626    proto_w: usize,
1627    y0: usize,
1628    x0: usize,
1629    roi_h: usize,
1630    roi_w: usize,
1631    num_protos: usize,
1632) -> crate::Result<ndarray::Array3<u8>> {
1633    use edgefirst_tensor::QuantMode;
1634    let stride_y = proto_w * num_protos;
1635
1636    // Precompute scaled coefficients + zp_offset (same as the old sigmoid kernel).
1637    let mut stack_scratch = [0.0_f32; 64];
1638    let mut heap_scratch: Vec<f32>;
1639    let scaled_coeff: &mut [f32] = if num_protos <= stack_scratch.len() {
1640        &mut stack_scratch[..num_protos]
1641    } else {
1642        heap_scratch = vec![0.0_f32; num_protos];
1643        heap_scratch.as_mut_slice()
1644    };
1645    let zp_offset: f32;
1646    match quant.mode() {
1647        QuantMode::PerTensorSymmetric { scale } => {
1648            for k in 0..num_protos {
1649                scaled_coeff[k] = coeff[k] * scale;
1650            }
1651            zp_offset = 0.0;
1652        }
1653        QuantMode::PerTensor { scale, zero_point } => {
1654            for k in 0..num_protos {
1655                scaled_coeff[k] = coeff[k] * scale;
1656            }
1657            zp_offset = zero_point as f32 * scaled_coeff.iter().take(num_protos).sum::<f32>();
1658        }
1659        QuantMode::PerChannelSymmetric { scales, axis } => {
1660            if axis != 2 {
1661                return Err(crate::Error::NotSupported(format!(
1662                    "per-channel quantization on axis {axis} not supported \
1663                     (only channel axis 2 is implemented on this kernel)"
1664                )));
1665            }
1666            for k in 0..num_protos {
1667                scaled_coeff[k] = coeff[k] * scales[k];
1668            }
1669            zp_offset = 0.0;
1670        }
1671        QuantMode::PerChannel {
1672            scales,
1673            zero_points,
1674            axis,
1675        } => {
1676            if axis != 2 {
1677                return Err(crate::Error::NotSupported(format!(
1678                    "per-channel quantization on axis {axis} not supported \
1679                     (only channel axis 2 is implemented on this kernel)"
1680                )));
1681            }
1682            for k in 0..num_protos {
1683                scaled_coeff[k] = coeff[k] * scales[k];
1684            }
1685            zp_offset = (0..num_protos)
1686                .map(|k| scaled_coeff[k] * zero_points[k] as f32)
1687                .sum();
1688        }
1689    }
1690
1691    let mut mask_buf = vec![0u8; roi_h * roi_w];
1692    for y in 0..roi_h {
1693        let row_base = (y0 + y) * stride_y + (x0) * num_protos;
1694        let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1695        for (x, out_px) in out_row.iter_mut().enumerate() {
1696            let base = row_base + x * num_protos;
1697            let mut acc = 0.0_f32;
1698            let mut k = 0;
1699            let chunks = num_protos / 4;
1700            for _ in 0..chunks {
1701                let p0 = protos[base + k] as f32;
1702                let p1 = protos[base + k + 1] as f32;
1703                let p2 = protos[base + k + 2] as f32;
1704                let p3 = protos[base + k + 3] as f32;
1705                acc += scaled_coeff[k] * p0
1706                    + scaled_coeff[k + 1] * p1
1707                    + scaled_coeff[k + 2] * p2
1708                    + scaled_coeff[k + 3] * p3;
1709                k += 4;
1710            }
1711            while k < num_protos {
1712                acc += scaled_coeff[k] * protos[base + k] as f32;
1713                k += 1;
1714            }
1715            if acc > zp_offset {
1716                *out_px = 255;
1717            }
1718        }
1719    }
1720    Ok(ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1721        .expect("mask_buf length matches roi_h * roi_w"))
1722}
1723
1724#[allow(clippy::too_many_arguments)]
1725fn scaled_segmentations_f32_slice(
1726    detect: &[crate::DetectBox],
1727    coeff_all: &[f32],
1728    protos: &[f32],
1729    proto_h: usize,
1730    proto_w: usize,
1731    num_protos: usize,
1732    letterbox: Option<[f32; 4]>,
1733    width: u32,
1734    height: u32,
1735) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
1736    scaled_run(
1737        detect,
1738        coeff_all,
1739        protos,
1740        proto_h,
1741        proto_w,
1742        num_protos,
1743        letterbox,
1744        width,
1745        height,
1746        1.0,
1747        |p, _| *p,
1748    )
1749}
1750
1751#[allow(clippy::too_many_arguments)]
1752fn scaled_segmentations_f16_slice(
1753    detect: &[crate::DetectBox],
1754    coeff_all: &[f32],
1755    protos: &[half::f16],
1756    proto_h: usize,
1757    proto_w: usize,
1758    num_protos: usize,
1759    letterbox: Option<[f32; 4]>,
1760    width: u32,
1761    height: u32,
1762) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
1763    scaled_run(
1764        detect,
1765        coeff_all,
1766        protos,
1767        proto_h,
1768        proto_w,
1769        num_protos,
1770        letterbox,
1771        width,
1772        height,
1773        1.0,
1774        |p: &half::f16, _| p.to_f32(),
1775    )
1776}
1777
1778#[allow(clippy::too_many_arguments)]
1779fn scaled_segmentations_i8_slice(
1780    detect: &[crate::DetectBox],
1781    coeff_all: &[f32],
1782    protos: &[i8],
1783    proto_h: usize,
1784    proto_w: usize,
1785    num_protos: usize,
1786    quant: &edgefirst_tensor::Quantization,
1787    letterbox: Option<[f32; 4]>,
1788    width: u32,
1789    height: u32,
1790) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
1791    use edgefirst_tensor::QuantMode;
1792    // Only per-tensor quantization supported on the scaled-path CPU kernel
1793    // today. Per-channel fits naturally into a future extension (would need
1794    // per-channel scaled coefficients in scaled_run's dot-product precompute).
1795    let (scale, zp) = match quant.mode() {
1796        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1797        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1798        QuantMode::PerChannel { axis, .. } | QuantMode::PerChannelSymmetric { axis, .. } => {
1799            return Err(crate::Error::NotSupported(format!(
1800                "per-channel quantization (axis={axis}) on scaled seg path \
1801                 not yet supported"
1802            )));
1803        }
1804    };
1805    scaled_run(
1806        detect,
1807        coeff_all,
1808        protos,
1809        proto_h,
1810        proto_w,
1811        num_protos,
1812        letterbox,
1813        width,
1814        height,
1815        scale,
1816        move |p: &i8, _| *p as f32 - zp,
1817    )
1818}
1819
1820// =============================================================================
1821// Integer-domain kernel: i8 coefficients × i8 protos → i32 → sign threshold.
1822//
1823// Eliminates all f32 conversion by working directly with raw quantized values.
1824// The math:
1825//   sign(dot(dequant(coeff), dequant(proto)))
1826//   = sign(Σ (c_raw - zp_c) · (p_raw - zp_p))
1827//   = sign(Σ c_raw·p_raw - zp_c·Σp_raw - zp_p·Σc_raw + N·zp_c·zp_p)
1828//   = sign(sdot(c_raw, p_raw) - zp_c·proto_sum[pixel] - bias_per_det)
1829//
1830// where bias_per_det = zp_p·Σc_raw - N·zp_c·zp_p  (precomputed once per det).
1831// =============================================================================
1832
1833/// Compute i8×i8 dot product (32 elements) → i32.
1834/// Platform-agnostic scalar fallback.
1835#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
1836#[inline(always)]
1837fn dot_i8_scalar(coeff: &[i8], proto: &[i8], n: usize) -> i32 {
1838    let mut acc: i32 = 0;
1839    let chunks = n / 4;
1840    let mut k = 0;
1841    for _ in 0..chunks {
1842        acc += coeff[k] as i32 * proto[k] as i32
1843            + coeff[k + 1] as i32 * proto[k + 1] as i32
1844            + coeff[k + 2] as i32 * proto[k + 2] as i32
1845            + coeff[k + 3] as i32 * proto[k + 3] as i32;
1846        k += 4;
1847    }
1848    while k < n {
1849        acc += coeff[k] as i32 * proto[k] as i32;
1850        k += 1;
1851    }
1852    acc
1853}
1854
1855/// NEON i8×i8→i32 dot product using smull+sadalp (works on ALL aarch64, A53+).
1856#[cfg(target_arch = "aarch64")]
1857#[inline(always)]
1858unsafe fn dot_i8_neon_base(coeff: *const i8, proto: *const i8, n: usize) -> i32 {
1859    use std::arch::aarch64::*;
1860    let mut acc = vdupq_n_s32(0);
1861    let full_chunks = n / 16;
1862    let mut offset = 0usize;
1863    for _ in 0..full_chunks {
1864        let c = vld1q_s8(coeff.add(offset));
1865        let p = vld1q_s8(proto.add(offset));
1866        // Widening multiply + pairwise accumulate (all aarch64).
1867        let lo = vmull_s8(vget_low_s8(c), vget_low_s8(p));
1868        let hi = vmull_high_s8(c, p);
1869        acc = vpadalq_s16(acc, lo);
1870        acc = vpadalq_s16(acc, hi);
1871        offset += 16;
1872    }
1873    // Handle remaining elements (for num_protos=32, full_chunks=2, remainder=0)
1874    let remainder = n - offset;
1875    if remainder >= 8 {
1876        let c = vld1_s8(coeff.add(offset));
1877        let p = vld1_s8(proto.add(offset));
1878        let prod = vmull_s8(c, p);
1879        acc = vpadalq_s16(acc, prod);
1880        offset += 8;
1881    }
1882    let mut scalar_acc = vaddvq_s32(acc);
1883    while offset < n {
1884        scalar_acc += *coeff.add(offset) as i32 * *proto.add(offset) as i32;
1885        offset += 1;
1886    }
1887    scalar_acc
1888}
1889
1890/// NEON i8×i8→i32 dot product using sdot (ARMv8.2-A dotprod, A55+).
1891/// Each `sdot` processes 16 i8 lanes → 4 i32 partial sums in one instruction,
1892/// replacing the 3-instruction smull+smull2+sadalp sequence.
1893#[cfg(target_arch = "aarch64")]
1894#[inline(always)]
1895unsafe fn dot_i8_neon_dotprod(coeff: *const i8, proto: *const i8, n: usize) -> i32 {
1896    use std::arch::aarch64::*;
1897    let mut acc = vdupq_n_s32(0);
1898    let full_chunks = n / 16;
1899    let mut offset = 0usize;
1900    for _ in 0..full_chunks {
1901        let c = vld1q_s8(coeff.add(offset));
1902        let p = vld1q_s8(proto.add(offset));
1903        // Enable dotprod extension locally so the assembler accepts sdot
1904        // even when compiling for baseline aarch64 (A53). At runtime we only
1905        // reach this path when HWCAP confirms dotprod support.
1906        let result: int32x4_t;
1907        core::arch::asm!(
1908            ".arch_extension dotprod",
1909            "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
1910            acc = inout(vreg) acc => result,
1911            a = in(vreg) c,
1912            b = in(vreg) p,
1913            options(pure, nomem, nostack),
1914        );
1915        acc = result;
1916        offset += 16;
1917    }
1918    let mut scalar_acc = vaddvq_s32(acc);
1919    // Tail: handle remainder (unlikely for num_protos=32, but correct)
1920    while offset < n {
1921        scalar_acc += *coeff.add(offset) as i32 * *proto.add(offset) as i32;
1922        offset += 1;
1923    }
1924    scalar_acc
1925}
1926
1927/// Compute i16×i8 dot product → i32. Platform-agnostic scalar fallback.
1928#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
1929#[inline(always)]
1930fn dot_i16_i8_scalar(coeff: &[i16], proto: &[i8], n: usize) -> i32 {
1931    let mut acc: i32 = 0;
1932    let chunks = n / 4;
1933    let mut k = 0;
1934    for _ in 0..chunks {
1935        acc += coeff[k] as i32 * proto[k] as i32
1936            + coeff[k + 1] as i32 * proto[k + 1] as i32
1937            + coeff[k + 2] as i32 * proto[k + 2] as i32
1938            + coeff[k + 3] as i32 * proto[k + 3] as i32;
1939        k += 4;
1940    }
1941    while k < n {
1942        acc += coeff[k] as i32 * proto[k] as i32;
1943        k += 1;
1944    }
1945    acc
1946}
1947
1948/// NEON i16×i8→i32 dot product using widening multiply-accumulate.
1949/// Processes 8 elements per iteration (vs 16 for i8×i8 dotprod).
1950#[cfg(target_arch = "aarch64")]
1951#[inline(always)]
1952unsafe fn dot_i16_i8_neon(coeff: *const i16, proto: *const i8, n: usize) -> i32 {
1953    use std::arch::aarch64::*;
1954    let mut acc = vdupq_n_s32(0);
1955    let full_chunks = n / 8;
1956    let mut offset = 0usize;
1957    for _ in 0..full_chunks {
1958        let c = vld1q_s16(coeff.add(offset));
1959        let p_raw = vld1_s8(proto.add(offset));
1960        let p = vmovl_s8(p_raw);
1961        acc = vmlal_s16(acc, vget_low_s16(c), vget_low_s16(p));
1962        acc = vmlal_high_s16(acc, c, p);
1963        offset += 8;
1964    }
1965    let mut scalar_acc = vaddvq_s32(acc);
1966    while offset < n {
1967        scalar_acc += *coeff.add(offset) as i32 * *proto.add(offset) as i32;
1968        offset += 1;
1969    }
1970    scalar_acc
1971}
1972
1973/// Compute the logit grid using the dotprod (sdot) path.
1974/// Separated into its own function so the compiler inlines the sdot asm fully.
1975#[cfg(target_arch = "aarch64")]
1976#[inline(always)]
1977#[allow(clippy::too_many_arguments)]
1978fn compute_logits_dotprod(
1979    logits: &mut [i32],
1980    coeff: &[i8],
1981    protos: &[i8],
1982    proto_sums: &[i32],
1983    proto_w: usize,
1984    proto_x0: usize,
1985    proto_y0: usize,
1986    roi_w: usize,
1987    roi_h: usize,
1988    stride_y: usize,
1989    num_protos: usize,
1990    zp_c: i32,
1991    bias: i32,
1992) {
1993    for ly_idx in 0..roi_h {
1994        let py = proto_y0 + ly_idx;
1995        let row_base = py * stride_y + proto_x0 * num_protos;
1996        for lx_idx in 0..roi_w {
1997            let pix_base = row_base + lx_idx * num_protos;
1998            let proto_px = &protos[pix_base..pix_base + num_protos];
1999            let raw_dot =
2000                unsafe { dot_i8_neon_dotprod(coeff.as_ptr(), proto_px.as_ptr(), num_protos) };
2001            let correction = if zp_c != 0 {
2002                zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2003            } else {
2004                0
2005            };
2006            logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2007        }
2008    }
2009}
2010
2011/// Compute the logit grid using the base NEON path (smull+sadalp).
2012/// Separated into its own function so the compiler inlines the NEON code fully.
2013#[cfg(target_arch = "aarch64")]
2014#[inline(always)]
2015#[allow(clippy::too_many_arguments)]
2016fn compute_logits_base(
2017    logits: &mut [i32],
2018    coeff: &[i8],
2019    protos: &[i8],
2020    proto_sums: &[i32],
2021    proto_w: usize,
2022    proto_x0: usize,
2023    proto_y0: usize,
2024    roi_w: usize,
2025    roi_h: usize,
2026    stride_y: usize,
2027    num_protos: usize,
2028    zp_c: i32,
2029    bias: i32,
2030) {
2031    for ly_idx in 0..roi_h {
2032        let py = proto_y0 + ly_idx;
2033        let row_base = py * stride_y + proto_x0 * num_protos;
2034        for lx_idx in 0..roi_w {
2035            let pix_base = row_base + lx_idx * num_protos;
2036            let proto_px = &protos[pix_base..pix_base + num_protos];
2037            let raw_dot =
2038                unsafe { dot_i8_neon_base(coeff.as_ptr(), proto_px.as_ptr(), num_protos) };
2039            let correction = if zp_c != 0 {
2040                zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2041            } else {
2042                0
2043            };
2044            logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2045        }
2046    }
2047}
2048
2049#[allow(clippy::too_many_arguments)]
2050fn scaled_segmentations_i8_i8(
2051    detect: &[crate::DetectBox],
2052    coeff_all: &[i8],
2053    coeff_quant: &edgefirst_tensor::Quantization,
2054    protos: &[i8],
2055    proto_quant: &edgefirst_tensor::Quantization,
2056    proto_h: usize,
2057    proto_w: usize,
2058    num_protos: usize,
2059    letterbox: Option<[f32; 4]>,
2060    width: u32,
2061    height: u32,
2062    layout: edgefirst_decoder::ProtoLayout,
2063) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
2064    use edgefirst_tensor::QuantMode;
2065
2066    let _span = tracing::trace_span!(
2067        "image.materialize_masks.kernel_i8_scaled",
2068        n = detect.len(),
2069        proto_h,
2070        proto_w,
2071        num_protos,
2072        width,
2073        height,
2074        ?layout,
2075    )
2076    .entered();
2077
2078    let zp_c: i32 = match coeff_quant.mode() {
2079        QuantMode::PerTensor { zero_point, .. } => zero_point,
2080        QuantMode::PerTensorSymmetric { .. } => 0,
2081        _ => {
2082            return Err(crate::Error::NotSupported(
2083                "per-channel coeff quantization not supported".into(),
2084            ))
2085        }
2086    };
2087    let zp_p: i32 = match proto_quant.mode() {
2088        QuantMode::PerTensor { zero_point, .. } => zero_point,
2089        QuantMode::PerTensorSymmetric { .. } => 0,
2090        _ => {
2091            return Err(crate::Error::NotSupported(
2092                "per-channel proto quantization not supported".into(),
2093            ))
2094        }
2095    };
2096
2097    let (lx0, lw, ly0, lh) = match letterbox {
2098        Some([lx0, ly0, lx1, ly1]) => {
2099            let lw = (lx1 - lx0).max(f32::EPSILON);
2100            let lh = (ly1 - ly0).max(f32::EPSILON);
2101            (lx0, lw, ly0, lh)
2102        }
2103        None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
2104    };
2105    let out_w = width as usize;
2106    let out_h = height as usize;
2107    let hw = proto_h * proto_w;
2108
2109    // Precompute proto_sum for the entire proto tensor (zero-point correction).
2110    let proto_sums: Vec<i32> = if zp_c != 0 {
2111        match layout {
2112            edgefirst_decoder::ProtoLayout::Nhwc => (0..hw)
2113                .map(|px_idx| {
2114                    let base = px_idx * num_protos;
2115                    let mut s: i32 = 0;
2116                    for k in 0..num_protos {
2117                        s += protos[base + k] as i32;
2118                    }
2119                    s
2120                })
2121                .collect(),
2122            edgefirst_decoder::ProtoLayout::Nchw => {
2123                let mut sums = vec![0i32; hw];
2124                for c in 0..num_protos {
2125                    let plane = &protos[c * hw..];
2126                    for (px, s) in sums.iter_mut().enumerate() {
2127                        *s += plane[px] as i32;
2128                    }
2129                }
2130                sums
2131            }
2132        }
2133    } else {
2134        Vec::new()
2135    };
2136
2137    // Detect dotprod support once, outside the hot loop.
2138    #[cfg(target_arch = "aarch64")]
2139    let use_dotprod = std::arch::is_aarch64_feature_detected!("dotprod");
2140
2141    // For NHWC layout, stride for row navigation.
2142    let stride_y = proto_w * num_protos;
2143
2144    detect
2145        .par_iter()
2146        .enumerate()
2147        .map(|(i, det)| {
2148            let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
2149            let bbox = det.bbox.to_canonical();
2150            let xmin = ((bbox.xmin - lx0) / lw).clamp(0.0, 1.0);
2151            let ymin = ((bbox.ymin - ly0) / lh).clamp(0.0, 1.0);
2152            let xmax = ((bbox.xmax - lx0) / lw).clamp(0.0, 1.0);
2153            let ymax = ((bbox.ymax - ly0) / lh).clamp(0.0, 1.0);
2154            let px0 = (xmin * out_w as f32).round() as usize;
2155            let py0 = (ymin * out_h as f32).round() as usize;
2156            let px1 = ((xmax * out_w as f32).round() as usize).min(out_w);
2157            let py1 = ((ymax * out_h as f32).round() as usize).min(out_h);
2158            let bbox_w = px1.saturating_sub(px0).max(1);
2159            let bbox_h = py1.saturating_sub(py0).max(1);
2160
2161            // Map output bbox → proto ROI.
2162            let sample_x_at = |px: f32| -> f32 {
2163                let model_x_norm = lx0 + (px + 0.5) / out_w as f32 * lw;
2164                model_x_norm * proto_w as f32 - 0.5
2165            };
2166            let sample_y_at = |py: f32| -> f32 {
2167                let model_y_norm = ly0 + (py + 0.5) / out_h as f32 * lh;
2168                model_y_norm * proto_h as f32 - 0.5
2169            };
2170            let s_x_min = sample_x_at(px0 as f32);
2171            let s_x_max = sample_x_at((px1 as f32) - 1.0);
2172            let s_y_min = sample_y_at(py0 as f32);
2173            let s_y_max = sample_y_at((py1 as f32) - 1.0);
2174            let proto_x0 = (s_x_min.floor() as isize)
2175                .max(0)
2176                .min(proto_w.saturating_sub(1) as isize) as usize;
2177            let proto_x1 = ((s_x_max.ceil() as isize) + 1).max(0).min(proto_w as isize) as usize;
2178            let proto_y0 = (s_y_min.floor() as isize)
2179                .max(0)
2180                .min(proto_h.saturating_sub(1) as isize) as usize;
2181            let proto_y1 = ((s_y_max.ceil() as isize) + 1).max(0).min(proto_h as isize) as usize;
2182            let roi_w = proto_x1.saturating_sub(proto_x0).max(1);
2183            let roi_h = proto_y1.saturating_sub(proto_y0).max(1);
2184
2185            // Per-detection bias.
2186            let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
2187            let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
2188
2189            // Step 2: Compute i32 logits at each proto-ROI pixel.
2190            let mut logits = vec![0_i32; roi_h * roi_w];
2191            match layout {
2192                edgefirst_decoder::ProtoLayout::Nhwc => {
2193                    #[cfg(target_arch = "aarch64")]
2194                    {
2195                        if use_dotprod {
2196                            compute_logits_dotprod(
2197                                &mut logits,
2198                                coeff,
2199                                protos,
2200                                &proto_sums,
2201                                proto_w,
2202                                proto_x0,
2203                                proto_y0,
2204                                roi_w,
2205                                roi_h,
2206                                stride_y,
2207                                num_protos,
2208                                zp_c,
2209                                bias,
2210                            );
2211                        } else {
2212                            compute_logits_base(
2213                                &mut logits,
2214                                coeff,
2215                                protos,
2216                                &proto_sums,
2217                                proto_w,
2218                                proto_x0,
2219                                proto_y0,
2220                                roi_w,
2221                                roi_h,
2222                                stride_y,
2223                                num_protos,
2224                                zp_c,
2225                                bias,
2226                            );
2227                        }
2228                    }
2229                    #[cfg(not(target_arch = "aarch64"))]
2230                    {
2231                        for ly_idx in 0..roi_h {
2232                            let py = proto_y0 + ly_idx;
2233                            let row_base = py * stride_y + proto_x0 * num_protos;
2234                            for lx_idx in 0..roi_w {
2235                                let pix_base = row_base + lx_idx * num_protos;
2236                                let proto_px = &protos[pix_base..pix_base + num_protos];
2237                                let raw_dot = dot_i8_scalar(coeff, proto_px, num_protos);
2238                                let correction = if zp_c != 0 {
2239                                    zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2240                                } else {
2241                                    0
2242                                };
2243                                logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2244                            }
2245                        }
2246                    }
2247                }
2248                edgefirst_decoder::ProtoLayout::Nchw => {
2249                    // Channel-major accumulation: contiguous reads per channel plane.
2250                    for c in 0..num_protos {
2251                        let plane = &protos[c * hw..];
2252                        let coeff_c = coeff[c] as i32;
2253                        for ly_idx in 0..roi_h {
2254                            let py = proto_y0 + ly_idx;
2255                            let row_start = py * proto_w + proto_x0;
2256                            let out_row_start = ly_idx * roi_w;
2257                            for lx_idx in 0..roi_w {
2258                                logits[out_row_start + lx_idx] +=
2259                                    coeff_c * plane[row_start + lx_idx] as i32;
2260                            }
2261                        }
2262                    }
2263                    // Apply zero-point correction and per-detection bias.
2264                    for ly_idx in 0..roi_h {
2265                        let py = proto_y0 + ly_idx;
2266                        for lx_idx in 0..roi_w {
2267                            let idx = ly_idx * roi_w + lx_idx;
2268                            let correction = if zp_c != 0 {
2269                                zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2270                            } else {
2271                                0
2272                            };
2273                            logits[idx] -= correction + bias;
2274                        }
2275                    }
2276                }
2277            }
2278
2279            // Step 3: Bilinear upsample i32 logits → binary mask with
2280            // sign-shortcut (skip interpolation when all 4 neighbors agree).
2281            let roi_last_x = roi_w.saturating_sub(1);
2282            let roi_last_y = roi_h.saturating_sub(1);
2283
2284            // X-coordinate LUT with fixed-point fraction (scale 1024).
2285            const FRAC_BITS: i32 = 10;
2286            const FRAC_SCALE: i32 = 1 << FRAC_BITS; // 1024
2287            let x_coords: Vec<(usize, usize, i32)> = (0..bbox_w)
2288                .map(|xi| {
2289                    let sample_x = sample_x_at((px0 + xi) as f32) - proto_x0 as f32;
2290                    let x_floor = sample_x.floor();
2291                    let x_lo = (x_floor as isize).max(0).min(roi_last_x as isize) as usize;
2292                    let x_hi = (x_lo + 1).min(roi_w - 1);
2293                    let x_frac = ((sample_x - x_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2294                    (x_lo, x_hi, x_frac)
2295                })
2296                .collect();
2297
2298            let mut tile_buf = vec![0u8; bbox_h * bbox_w];
2299            for yi in 0..bbox_h {
2300                let sample_y = sample_y_at((py0 + yi) as f32) - proto_y0 as f32;
2301                let y_floor = sample_y.floor();
2302                let y_lo = (y_floor as isize).max(0).min(roi_last_y as isize) as usize;
2303                let y_hi = (y_lo + 1).min(roi_h - 1);
2304                let y_frac = ((sample_y - y_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2305                let y_frac_inv = FRAC_SCALE - y_frac;
2306                let row_lo = &logits[y_lo * roi_w..y_lo * roi_w + roi_w];
2307                let row_hi = &logits[y_hi * roi_w..y_hi * roi_w + roi_w];
2308                let out_row = &mut tile_buf[yi * bbox_w..(yi + 1) * bbox_w];
2309
2310                for (xi, &(x_lo, x_hi, x_frac)) in x_coords.iter().enumerate() {
2311                    let tl = row_lo[x_lo];
2312                    let tr = row_lo[x_hi];
2313                    let bl = row_hi[x_lo];
2314                    let br = row_hi[x_hi];
2315
2316                    // Sign-shortcut: if all 4 corners have the same sign,
2317                    // the bilinear interpolation (positive-weight combination)
2318                    // preserves that sign. Skip arithmetic for ~80% of pixels.
2319                    if (tl & tr & bl & br) < 0 {
2320                        // All negative → output 0 (already zero).
2321                        continue;
2322                    }
2323                    if tl > 0 && tr > 0 && bl > 0 && br > 0 {
2324                        // All strictly positive → output 255.
2325                        out_row[xi] = 255;
2326                        continue;
2327                    }
2328
2329                    // Boundary pixel: fixed-point bilinear in i64.
2330                    let x_frac_inv = FRAC_SCALE - x_frac;
2331                    let l0 = tl as i64 * x_frac_inv as i64 + tr as i64 * x_frac as i64;
2332                    let l1 = bl as i64 * x_frac_inv as i64 + br as i64 * x_frac as i64;
2333                    let logit = l0 * y_frac_inv as i64 + l1 * y_frac as i64;
2334                    out_row[xi] = if logit > 0 { 255 } else { 0 };
2335                }
2336            }
2337
2338            let tile = ndarray::Array3::from_shape_vec((bbox_h, bbox_w, 1), tile_buf)
2339                .expect("tile_buf length matches bbox_h * bbox_w");
2340            Ok(edgefirst_decoder::Segmentation {
2341                xmin,
2342                ymin,
2343                xmax,
2344                ymax,
2345                segmentation: tile,
2346            })
2347        })
2348        .collect()
2349}
2350
2351#[allow(clippy::too_many_arguments)]
2352fn scaled_segmentations_i16_i8(
2353    detect: &[crate::DetectBox],
2354    coeff_all: &[i16],
2355    coeff_quant: &edgefirst_tensor::Quantization,
2356    protos: &[i8],
2357    proto_quant: &edgefirst_tensor::Quantization,
2358    proto_h: usize,
2359    proto_w: usize,
2360    num_protos: usize,
2361    letterbox: Option<[f32; 4]>,
2362    width: u32,
2363    height: u32,
2364    layout: edgefirst_decoder::ProtoLayout,
2365) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
2366    use edgefirst_tensor::QuantMode;
2367
2368    let _span = tracing::trace_span!(
2369        "image.materialize_masks.kernel_i16xi8_scaled",
2370        n = detect.len(),
2371        proto_h,
2372        proto_w,
2373        num_protos,
2374        width,
2375        height,
2376        ?layout,
2377    )
2378    .entered();
2379
2380    let zp_c: i32 = match coeff_quant.mode() {
2381        QuantMode::PerTensor { zero_point, .. } => zero_point,
2382        QuantMode::PerTensorSymmetric { .. } => 0,
2383        _ => {
2384            return Err(crate::Error::NotSupported(
2385                "per-channel coeff quantization not supported".into(),
2386            ))
2387        }
2388    };
2389    let zp_p: i32 = match proto_quant.mode() {
2390        QuantMode::PerTensor { zero_point, .. } => zero_point,
2391        QuantMode::PerTensorSymmetric { .. } => 0,
2392        _ => {
2393            return Err(crate::Error::NotSupported(
2394                "per-channel proto quantization not supported".into(),
2395            ))
2396        }
2397    };
2398
2399    let (lx0, lw, ly0, lh) = match letterbox {
2400        Some([lx0, ly0, lx1, ly1]) => {
2401            let lw = (lx1 - lx0).max(f32::EPSILON);
2402            let lh = (ly1 - ly0).max(f32::EPSILON);
2403            (lx0, lw, ly0, lh)
2404        }
2405        None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
2406    };
2407    let out_w = width as usize;
2408    let out_h = height as usize;
2409    let hw = proto_h * proto_w;
2410
2411    // Precompute proto_sum for the entire proto tensor (zero-point correction).
2412    let proto_sums: Vec<i32> = if zp_c != 0 {
2413        match layout {
2414            edgefirst_decoder::ProtoLayout::Nhwc => (0..hw)
2415                .map(|px_idx| {
2416                    let base = px_idx * num_protos;
2417                    let mut s: i32 = 0;
2418                    for k in 0..num_protos {
2419                        s += protos[base + k] as i32;
2420                    }
2421                    s
2422                })
2423                .collect(),
2424            edgefirst_decoder::ProtoLayout::Nchw => {
2425                let mut sums = vec![0i32; hw];
2426                for c in 0..num_protos {
2427                    let plane = &protos[c * hw..];
2428                    for (px, s) in sums.iter_mut().enumerate() {
2429                        *s += plane[px] as i32;
2430                    }
2431                }
2432                sums
2433            }
2434        }
2435    } else {
2436        Vec::new()
2437    };
2438
2439    // For NHWC layout, stride for row navigation.
2440    let stride_y = proto_w * num_protos;
2441
2442    detect
2443        .par_iter()
2444        .enumerate()
2445        .map(|(i, det)| {
2446            let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
2447            let bbox = det.bbox.to_canonical();
2448            let xmin = ((bbox.xmin - lx0) / lw).clamp(0.0, 1.0);
2449            let ymin = ((bbox.ymin - ly0) / lh).clamp(0.0, 1.0);
2450            let xmax = ((bbox.xmax - lx0) / lw).clamp(0.0, 1.0);
2451            let ymax = ((bbox.ymax - ly0) / lh).clamp(0.0, 1.0);
2452            let px0 = (xmin * out_w as f32).round() as usize;
2453            let py0 = (ymin * out_h as f32).round() as usize;
2454            let px1 = ((xmax * out_w as f32).round() as usize).min(out_w);
2455            let py1 = ((ymax * out_h as f32).round() as usize).min(out_h);
2456            let bbox_w = px1.saturating_sub(px0).max(1);
2457            let bbox_h = py1.saturating_sub(py0).max(1);
2458
2459            // Map output bbox → proto ROI.
2460            let sample_x_at = |px: f32| -> f32 {
2461                let model_x_norm = lx0 + (px + 0.5) / out_w as f32 * lw;
2462                model_x_norm * proto_w as f32 - 0.5
2463            };
2464            let sample_y_at = |py: f32| -> f32 {
2465                let model_y_norm = ly0 + (py + 0.5) / out_h as f32 * lh;
2466                model_y_norm * proto_h as f32 - 0.5
2467            };
2468            let s_x_min = sample_x_at(px0 as f32);
2469            let s_x_max = sample_x_at((px1 as f32) - 1.0);
2470            let s_y_min = sample_y_at(py0 as f32);
2471            let s_y_max = sample_y_at((py1 as f32) - 1.0);
2472            let proto_x0 = (s_x_min.floor() as isize)
2473                .max(0)
2474                .min(proto_w.saturating_sub(1) as isize) as usize;
2475            let proto_x1 = ((s_x_max.ceil() as isize) + 1).max(0).min(proto_w as isize) as usize;
2476            let proto_y0 = (s_y_min.floor() as isize)
2477                .max(0)
2478                .min(proto_h.saturating_sub(1) as isize) as usize;
2479            let proto_y1 = ((s_y_max.ceil() as isize) + 1).max(0).min(proto_h as isize) as usize;
2480            let roi_w = proto_x1.saturating_sub(proto_x0).max(1);
2481            let roi_h = proto_y1.saturating_sub(proto_y0).max(1);
2482
2483            // Per-detection bias.
2484            let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
2485            let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
2486
2487            // Step 2: Compute i32 logits at each proto-ROI pixel.
2488            let mut logits = vec![0_i32; roi_h * roi_w];
2489            match layout {
2490                edgefirst_decoder::ProtoLayout::Nhwc => {
2491                    #[cfg(target_arch = "aarch64")]
2492                    {
2493                        for ly_idx in 0..roi_h {
2494                            let py = proto_y0 + ly_idx;
2495                            let row_base = py * stride_y + proto_x0 * num_protos;
2496                            for lx_idx in 0..roi_w {
2497                                let pix_base = row_base + lx_idx * num_protos;
2498                                let proto_px = &protos[pix_base..pix_base + num_protos];
2499                                let raw_dot = unsafe {
2500                                    dot_i16_i8_neon(coeff.as_ptr(), proto_px.as_ptr(), num_protos)
2501                                };
2502                                let correction = if zp_c != 0 {
2503                                    zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2504                                } else {
2505                                    0
2506                                };
2507                                logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2508                            }
2509                        }
2510                    }
2511                    #[cfg(not(target_arch = "aarch64"))]
2512                    {
2513                        for ly_idx in 0..roi_h {
2514                            let py = proto_y0 + ly_idx;
2515                            let row_base = py * stride_y + proto_x0 * num_protos;
2516                            for lx_idx in 0..roi_w {
2517                                let pix_base = row_base + lx_idx * num_protos;
2518                                let proto_px = &protos[pix_base..pix_base + num_protos];
2519                                let raw_dot = dot_i16_i8_scalar(coeff, proto_px, num_protos);
2520                                let correction = if zp_c != 0 {
2521                                    zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2522                                } else {
2523                                    0
2524                                };
2525                                logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2526                            }
2527                        }
2528                    }
2529                }
2530                edgefirst_decoder::ProtoLayout::Nchw => {
2531                    // Channel-major accumulation: contiguous reads per channel plane.
2532                    for c in 0..num_protos {
2533                        let plane = &protos[c * hw..];
2534                        let coeff_c = coeff[c] as i32;
2535                        for ly_idx in 0..roi_h {
2536                            let py = proto_y0 + ly_idx;
2537                            let row_start = py * proto_w + proto_x0;
2538                            let out_row_start = ly_idx * roi_w;
2539                            for lx_idx in 0..roi_w {
2540                                logits[out_row_start + lx_idx] +=
2541                                    coeff_c * plane[row_start + lx_idx] as i32;
2542                            }
2543                        }
2544                    }
2545                    // Apply zero-point correction and per-detection bias.
2546                    for ly_idx in 0..roi_h {
2547                        let py = proto_y0 + ly_idx;
2548                        for lx_idx in 0..roi_w {
2549                            let idx = ly_idx * roi_w + lx_idx;
2550                            let correction = if zp_c != 0 {
2551                                zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2552                            } else {
2553                                0
2554                            };
2555                            logits[idx] -= correction + bias;
2556                        }
2557                    }
2558                }
2559            }
2560
2561            // Step 3: Bilinear upsample i32 logits → binary mask with
2562            // sign-shortcut (skip interpolation when all 4 neighbors agree).
2563            let roi_last_x = roi_w.saturating_sub(1);
2564            let roi_last_y = roi_h.saturating_sub(1);
2565
2566            // X-coordinate LUT with fixed-point fraction (scale 1024).
2567            const FRAC_BITS: i32 = 10;
2568            const FRAC_SCALE: i32 = 1 << FRAC_BITS; // 1024
2569            let x_coords: Vec<(usize, usize, i32)> = (0..bbox_w)
2570                .map(|xi| {
2571                    let sample_x = sample_x_at((px0 + xi) as f32) - proto_x0 as f32;
2572                    let x_floor = sample_x.floor();
2573                    let x_lo = (x_floor as isize).max(0).min(roi_last_x as isize) as usize;
2574                    let x_hi = (x_lo + 1).min(roi_w - 1);
2575                    let x_frac = ((sample_x - x_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2576                    (x_lo, x_hi, x_frac)
2577                })
2578                .collect();
2579
2580            let mut tile_buf = vec![0u8; bbox_h * bbox_w];
2581            for yi in 0..bbox_h {
2582                let sample_y = sample_y_at((py0 + yi) as f32) - proto_y0 as f32;
2583                let y_floor = sample_y.floor();
2584                let y_lo = (y_floor as isize).max(0).min(roi_last_y as isize) as usize;
2585                let y_hi = (y_lo + 1).min(roi_h - 1);
2586                let y_frac = ((sample_y - y_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2587                let y_frac_inv = FRAC_SCALE - y_frac;
2588                let row_lo = &logits[y_lo * roi_w..y_lo * roi_w + roi_w];
2589                let row_hi = &logits[y_hi * roi_w..y_hi * roi_w + roi_w];
2590                let out_row = &mut tile_buf[yi * bbox_w..(yi + 1) * bbox_w];
2591
2592                for (xi, &(x_lo, x_hi, x_frac)) in x_coords.iter().enumerate() {
2593                    let tl = row_lo[x_lo];
2594                    let tr = row_lo[x_hi];
2595                    let bl = row_hi[x_lo];
2596                    let br = row_hi[x_hi];
2597
2598                    // Sign-shortcut: if all 4 corners have the same sign,
2599                    // the bilinear interpolation (positive-weight combination)
2600                    // preserves that sign. Skip arithmetic for ~80% of pixels.
2601                    if (tl & tr & bl & br) < 0 {
2602                        // All negative → output 0 (already zero).
2603                        continue;
2604                    }
2605                    if tl > 0 && tr > 0 && bl > 0 && br > 0 {
2606                        // All strictly positive → output 255.
2607                        out_row[xi] = 255;
2608                        continue;
2609                    }
2610
2611                    // Boundary pixel: fixed-point bilinear in i64.
2612                    let x_frac_inv = FRAC_SCALE - x_frac;
2613                    let l0 = tl as i64 * x_frac_inv as i64 + tr as i64 * x_frac as i64;
2614                    let l1 = bl as i64 * x_frac_inv as i64 + br as i64 * x_frac as i64;
2615                    let logit = l0 * y_frac_inv as i64 + l1 * y_frac as i64;
2616                    out_row[xi] = if logit > 0 { 255 } else { 0 };
2617                }
2618            }
2619
2620            let tile = ndarray::Array3::from_shape_vec((bbox_h, bbox_w, 1), tile_buf)
2621                .expect("tile_buf length matches bbox_h * bbox_w");
2622            Ok(edgefirst_decoder::Segmentation {
2623                xmin,
2624                ymin,
2625                xmax,
2626                ymax,
2627                segmentation: tile,
2628            })
2629        })
2630        .collect()
2631}
2632
2633#[allow(clippy::too_many_arguments)]
2634fn scaled_run<P: Copy + Sync>(
2635    detect: &[crate::DetectBox],
2636    coeff_all: &[f32],
2637    protos: &[P],
2638    proto_h: usize,
2639    proto_w: usize,
2640    num_protos: usize,
2641    letterbox: Option<[f32; 4]>,
2642    width: u32,
2643    height: u32,
2644    acc_scale: f32,
2645    load_f32: impl Fn(&P, f32) -> f32 + Copy + Sync,
2646) -> crate::Result<Vec<edgefirst_decoder::Segmentation>> {
2647    let (lx0, lw, ly0, lh) = match letterbox {
2648        Some([lx0, ly0, lx1, ly1]) => {
2649            let lw = (lx1 - lx0).max(f32::EPSILON);
2650            let lh = (ly1 - ly0).max(f32::EPSILON);
2651            (lx0, lw, ly0, lh)
2652        }
2653        None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
2654    };
2655    let out_w = width as usize;
2656    let out_h = height as usize;
2657    let stride_y = proto_w * num_protos;
2658
2659    // Parallelise across detections. Each detection produces an
2660    // independent ndarray::Array3<u8> tile from a read-only proto slice +
2661    // its own coeff slice; no shared mutable state.
2662    //
2663    // Algorithm (restores the spirit of PR #54's batched-GEMM optimisation
2664    // that PR #51's f16 dispatch refactor inadvertently removed):
2665    //
2666    //   1. Map the output bbox back to a proto-plane ROI (with 1-px margin
2667    //      so the bilinear sampling at the output edges has neighbours).
2668    //   2. Precompute *f32 logits* at every proto pixel inside that ROI by
2669    //      doing a single K-wide dot product per proto pixel — once, not
2670    //      once per output pixel.
2671    //   3. For each output pixel, bilinear-interpolate the scalar f32 logit
2672    //      from the 4 surrounding proto-roi pixels, apply sigmoid, and
2673    //      threshold to {0, 255}.
2674    //
2675    // For typical YOLO-seg: proto_roi ~ 30×30 = 900 px × K=32 = 28.8K dot
2676    // ops vs the legacy "bilinear sample then dot at every output pixel"
2677    // which costs bbox_h × bbox_w × 4 × K = ~1.3M ops at 100×100 output
2678    // bbox. ~45× fewer FMAs at this size; the bilinear upsample of a
2679    // scalar plane (no inner K loop) is comparatively negligible.
2680    detect
2681        .par_iter()
2682        .enumerate()
2683        .map(|(i, det)| {
2684            let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
2685            let bbox = det.bbox.to_canonical();
2686            let xmin = ((bbox.xmin - lx0) / lw).clamp(0.0, 1.0);
2687            let ymin = ((bbox.ymin - ly0) / lh).clamp(0.0, 1.0);
2688            let xmax = ((bbox.xmax - lx0) / lw).clamp(0.0, 1.0);
2689            let ymax = ((bbox.ymax - ly0) / lh).clamp(0.0, 1.0);
2690            let px0 = (xmin * out_w as f32).round() as usize;
2691            let py0 = (ymin * out_h as f32).round() as usize;
2692            let px1 = ((xmax * out_w as f32).round() as usize).min(out_w);
2693            let py1 = ((ymax * out_h as f32).round() as usize).min(out_h);
2694            let bbox_w = px1.saturating_sub(px0).max(1);
2695            let bbox_h = py1.saturating_sub(py0).max(1);
2696
2697            // Step 1 — proto-plane ROI for this detection's output bbox.
2698            // Map the four output bbox corners back to proto coords and
2699            // expand by 1 pixel in each direction so the bilinear sampler
2700            // at the bbox boundary has both neighbours.
2701            let sample_x_at = |px: f32| -> f32 {
2702                let model_x_norm = lx0 + (px + 0.5) / out_w as f32 * lw;
2703                model_x_norm * proto_w as f32 - 0.5
2704            };
2705            let sample_y_at = |py: f32| -> f32 {
2706                let model_y_norm = ly0 + (py + 0.5) / out_h as f32 * lh;
2707                model_y_norm * proto_h as f32 - 0.5
2708            };
2709            let s_x_min = sample_x_at(px0 as f32);
2710            let s_x_max = sample_x_at((px1 as f32) - 1.0);
2711            let s_y_min = sample_y_at(py0 as f32);
2712            let s_y_max = sample_y_at((py1 as f32) - 1.0);
2713            // Floor min, ceil max+1 to include both bilinear neighbours.
2714            // Start indices are used as direct bases into `protos`, so clamp
2715            // them to the last valid index, not to the exclusive upper bound.
2716            let proto_x0 = (s_x_min.floor() as isize)
2717                .max(0)
2718                .min(proto_w.saturating_sub(1) as isize) as usize;
2719            let proto_x1 = ((s_x_max.ceil() as isize) + 1).max(0).min(proto_w as isize) as usize;
2720            let proto_y0 = (s_y_min.floor() as isize)
2721                .max(0)
2722                .min(proto_h.saturating_sub(1) as isize) as usize;
2723            let proto_y1 = ((s_y_max.ceil() as isize) + 1).max(0).min(proto_h as isize) as usize;
2724            let roi_w = proto_x1.saturating_sub(proto_x0).max(1);
2725            let roi_h = proto_y1.saturating_sub(proto_y0).max(1);
2726
2727            // Step 2 — precompute f32 logits at every proto-roi pixel.
2728            // logits[(py - proto_y0) * roi_w + (px - proto_x0)] = dot(coeff, proto[py, px, :])
2729            //
2730            // Since the final threshold is `logit > 0` (O1) and bilinear
2731            // interpolation is a positive-weight linear combination,
2732            // `acc_scale * interp(logits) > 0 ⟺ interp(logits) > 0` when
2733            // acc_scale > 0. We therefore skip the per-pixel `acc_scale *`
2734            // multiply entirely, storing raw dot products.
2735            if !acc_scale.is_finite() || acc_scale <= 0.0 {
2736                return Err(crate::Error::NotSupported(format!(
2737                    "acc_scale must be finite and positive for sign-threshold optimization (got {acc_scale})"
2738                )));
2739            }
2740            let _ = acc_scale; // Scale-invariant: only sign matters.
2741            let mut logits = vec![0.0_f32; roi_h * roi_w];
2742            for ly_idx in 0..roi_h {
2743                let py = proto_y0 + ly_idx;
2744                let row_base = py * stride_y + proto_x0 * num_protos;
2745                for lx_idx in 0..roi_w {
2746                    let pix_base = row_base + lx_idx * num_protos;
2747                    let mut acc = 0.0_f32;
2748                    // 4-wide unroll to help auto-vectorization.
2749                    let mut k = 0;
2750                    let chunks = num_protos / 4;
2751                    for _ in 0..chunks {
2752                        acc += coeff[k] * load_f32(&protos[pix_base + k], 0.0)
2753                            + coeff[k + 1] * load_f32(&protos[pix_base + k + 1], 0.0)
2754                            + coeff[k + 2] * load_f32(&protos[pix_base + k + 2], 0.0)
2755                            + coeff[k + 3] * load_f32(&protos[pix_base + k + 3], 0.0);
2756                        k += 4;
2757                    }
2758                    while k < num_protos {
2759                        acc += coeff[k] * load_f32(&protos[pix_base + k], 0.0);
2760                        k += 1;
2761                    }
2762                    logits[ly_idx * roi_w + lx_idx] = acc;
2763                }
2764            }
2765
2766            // Step 3 — bilinear upsample logits → binary mask.
2767            //
2768            // O1: sigmoid(x) > 0.5 ⟺ x > 0 (sigmoid is strictly monotonic,
2769            // and acc_scale > 0 preserves sign). The sign threshold replaces
2770            // the old fast_sigmoid approximation, saving ~15 cycles/pixel.
2771            //
2772            // O5: Pre-compute bilinear sample coordinates. sample_x_at /
2773            // sample_y_at depend only on pixel index, not on logit values.
2774            // Building lookup tables avoids redundant float ops in the inner
2775            // loop (floor, clamp, isize cast per pixel).
2776            let roi_last_x = roi_w.saturating_sub(1);
2777            let roi_last_y = roi_h.saturating_sub(1);
2778
2779            // X-coordinate LUT (shared across all rows).
2780            let x_coords: Vec<(u32, u32, f32)> = (0..bbox_w)
2781                .map(|xi| {
2782                    let sample_x = sample_x_at((px0 + xi) as f32) - proto_x0 as f32;
2783                    let x_floor = sample_x.floor();
2784                    let x_lo = (x_floor as isize).max(0).min(roi_last_x as isize) as u32;
2785                    let x_hi = (x_lo as usize + 1).min(roi_w - 1) as u32;
2786                    let x_frac = (sample_x - x_floor).clamp(0.0, 1.0);
2787                    (x_lo, x_hi, x_frac)
2788                })
2789                .collect();
2790
2791            // Write the output tile through a contiguous slice to avoid
2792            // ndarray's per-element bounds checks + stride arithmetic.
2793            let mut tile_buf = vec![0u8; bbox_h * bbox_w];
2794            for yi in 0..bbox_h {
2795                let sample_y = sample_y_at((py0 + yi) as f32) - proto_y0 as f32;
2796                let y_floor = sample_y.floor();
2797                let y_lo = (y_floor as isize).max(0).min(roi_last_y as isize) as usize;
2798                let y_hi = (y_lo + 1).min(roi_h - 1);
2799                let y_frac = (sample_y - y_floor).clamp(0.0, 1.0);
2800                let row_lo = &logits[y_lo * roi_w..y_lo * roi_w + roi_w];
2801                let row_hi = &logits[y_hi * roi_w..y_hi * roi_w + roi_w];
2802                let out_row = &mut tile_buf[yi * bbox_w..(yi + 1) * bbox_w];
2803                for (xi, &(x_lo, x_hi, x_frac)) in x_coords.iter().enumerate() {
2804                    let (xl, xh) = (x_lo as usize, x_hi as usize);
2805                    let l0 = row_lo[xl] + (row_lo[xh] - row_lo[xl]) * x_frac;
2806                    let l1 = row_hi[xl] + (row_hi[xh] - row_hi[xl]) * x_frac;
2807                    let logit = l0 + (l1 - l0) * y_frac;
2808                    out_row[xi] = if logit > 0.0 { 255 } else { 0 };
2809                }
2810            }
2811            // Wrap into the expected Array3<u8> shape [bbox_h, bbox_w, 1].
2812            let tile = ndarray::Array3::from_shape_vec((bbox_h, bbox_w, 1), tile_buf)
2813                .expect("tile_buf length matches bbox_h * bbox_w");
2814            Ok(edgefirst_decoder::Segmentation {
2815                xmin,
2816                ymin,
2817                xmax,
2818                ymax,
2819                segmentation: tile,
2820            })
2821        })
2822        .collect()
2823}
2824
2825#[cfg(test)]
2826mod tests {
2827    use super::CPUProcessor;
2828    use edgefirst_decoder::{BoundingBox, DetectBox, ProtoData, ProtoLayout};
2829    use edgefirst_tensor::{Quantization, Tensor, TensorDyn};
2830
2831    const PROTO_H: usize = 4;
2832    const PROTO_W: usize = 4;
2833    const NUM_PROTOS: usize = 8;
2834
2835    fn det(xmin: f32, ymin: f32, xmax: f32, ymax: f32) -> DetectBox {
2836        DetectBox {
2837            bbox: BoundingBox {
2838                xmin,
2839                ymin,
2840                xmax,
2841                ymax,
2842            },
2843            score: 0.9,
2844            label: 0,
2845        }
2846    }
2847
2848    fn make_i8_quant(shape: &[usize], data: &[i8], scale: f32, zp: i32) -> TensorDyn {
2849        let t = Tensor::<i8>::from_slice(data, shape).unwrap();
2850        let t = t
2851            .with_quantization(Quantization::per_tensor(scale, zp))
2852            .unwrap();
2853        TensorDyn::I8(t)
2854    }
2855
2856    fn make_i16_quant(shape: &[usize], data: &[i16], scale: f32, zp: i32) -> TensorDyn {
2857        let t = Tensor::<i16>::from_slice(data, shape).unwrap();
2858        let t = t
2859            .with_quantization(Quantization::per_tensor(scale, zp))
2860            .unwrap();
2861        TensorDyn::I16(t)
2862    }
2863
2864    fn make_i16_raw(shape: &[usize], data: &[i16]) -> TensorDyn {
2865        let t = Tensor::<i16>::from_slice(data, shape).unwrap();
2866        TensorDyn::I16(t)
2867    }
2868
2869    fn make_f32(shape: &[usize], data: &[f32]) -> TensorDyn {
2870        let t = Tensor::<f32>::from_slice(data, shape).unwrap();
2871        TensorDyn::F32(t)
2872    }
2873
2874    fn gen_protos_i8(h: usize, w: usize, k: usize) -> Vec<i8> {
2875        (0..h * w * k).map(|i| (i % 127) as i8).collect()
2876    }
2877
2878    fn gen_coeffs_i16(n: usize, k: usize) -> Vec<i16> {
2879        (0..n * k)
2880            .map(|i| ((i as i32 % 201) - 100) as i16)
2881            .collect()
2882    }
2883
2884    fn gen_coeffs_i8(n: usize, k: usize) -> Vec<i8> {
2885        (0..n * k).map(|i| ((i as i32 % 201) - 100) as i8).collect()
2886    }
2887
2888    // ── Proto-resolution: i16×i8 fast path (quantized) ─────────────
2889
2890    #[test]
2891    fn materialize_proto_i16_i8_quant_produces_masks() {
2892        let cpu = CPUProcessor::new();
2893        let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
2894        let protos = make_i8_quant(
2895            &[PROTO_H, PROTO_W, NUM_PROTOS],
2896            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
2897            0.02,
2898            0,
2899        );
2900        let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 0);
2901        let proto_data = ProtoData {
2902            mask_coefficients: coeffs,
2903            protos,
2904            layout: ProtoLayout::Nhwc,
2905        };
2906        let result = cpu.materialize_segmentations(&detect, &proto_data, None);
2907        assert!(result.is_ok(), "materialize failed: {:?}", result.err());
2908        let segs = result.unwrap();
2909        assert_eq!(segs.len(), 1);
2910        let seg = &segs[0];
2911        assert!(seg.segmentation.shape()[0] > 0);
2912        assert!(seg.segmentation.shape()[1] > 0);
2913    }
2914
2915    // ── Proto-resolution: i16 missing quant → f32 fallback ─────────
2916
2917    #[test]
2918    fn materialize_proto_i16_no_quant_falls_back_to_f32() {
2919        let cpu = CPUProcessor::new();
2920        let detect = vec![det(0.2, 0.2, 0.8, 0.8)];
2921        let protos = make_i8_quant(
2922            &[PROTO_H, PROTO_W, NUM_PROTOS],
2923            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
2924            0.02,
2925            0,
2926        );
2927        // I16 coefficients WITHOUT quantization — fast path should be
2928        // skipped, f32 fallback should widen raw i16 values.
2929        let coeffs = make_i16_raw(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS));
2930        let proto_data = ProtoData {
2931            mask_coefficients: coeffs,
2932            protos,
2933            layout: ProtoLayout::Nhwc,
2934        };
2935        let result = cpu.materialize_segmentations(&detect, &proto_data, None);
2936        assert!(
2937            result.is_ok(),
2938            "missing coeff quant should fall back to f32 path, got: {:?}",
2939            result.err()
2940        );
2941        assert_eq!(result.unwrap().len(), 1);
2942    }
2943
2944    // ── Scaled: i16×i8 fast path (quantized) ───────────────────────
2945
2946    #[test]
2947    fn materialize_scaled_i16_i8_quant_produces_masks() {
2948        let cpu = CPUProcessor::new();
2949        let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
2950        let protos = make_i8_quant(
2951            &[PROTO_H, PROTO_W, NUM_PROTOS],
2952            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
2953            0.02,
2954            0,
2955        );
2956        let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 0);
2957        let proto_data = ProtoData {
2958            mask_coefficients: coeffs,
2959            protos,
2960            layout: ProtoLayout::Nhwc,
2961        };
2962        let result = cpu.materialize_scaled_segmentations(&detect, &proto_data, None, 64, 64);
2963        assert!(
2964            result.is_ok(),
2965            "materialize_scaled failed: {:?}",
2966            result.err()
2967        );
2968        let segs = result.unwrap();
2969        assert_eq!(segs.len(), 1);
2970        let seg = &segs[0];
2971        assert!(seg.segmentation.shape()[0] > 0);
2972        assert!(seg.segmentation.shape()[1] > 0);
2973    }
2974
2975    // ── Scaled: i16 missing quant → f32 fallback ───────────────────
2976
2977    #[test]
2978    fn materialize_scaled_i16_no_quant_falls_back_to_f32() {
2979        let cpu = CPUProcessor::new();
2980        let detect = vec![det(0.2, 0.2, 0.8, 0.8)];
2981        let protos = make_i8_quant(
2982            &[PROTO_H, PROTO_W, NUM_PROTOS],
2983            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
2984            0.02,
2985            0,
2986        );
2987        let coeffs = make_i16_raw(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS));
2988        let proto_data = ProtoData {
2989            mask_coefficients: coeffs,
2990            protos,
2991            layout: ProtoLayout::Nhwc,
2992        };
2993        let result = cpu.materialize_scaled_segmentations(&detect, &proto_data, None, 64, 64);
2994        assert!(
2995            result.is_ok(),
2996            "missing coeff quant should fall back to f32 path, got: {:?}",
2997            result.err()
2998        );
2999        assert_eq!(result.unwrap().len(), 1);
3000    }
3001
3002    // ── i16×i8 parity with f32 reference ───────────────────────────
3003
3004    #[test]
3005    fn materialize_proto_i16_i8_matches_f32_reference() {
3006        let cpu = CPUProcessor::new();
3007        let detect = vec![det(0.1, 0.1, 0.9, 0.9), det(0.3, 0.3, 0.7, 0.7)];
3008        let n_det = detect.len();
3009        let scale_c = 0.01_f32;
3010        let scale_p = 0.02_f32;
3011        let raw_protos = gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS);
3012        let raw_coeffs = gen_coeffs_i16(n_det, NUM_PROTOS);
3013
3014        // Build f32 reference (dequantized manually).
3015        let protos_f32: Vec<f32> = raw_protos.iter().map(|&v| v as f32 * scale_p).collect();
3016        let coeffs_f32: Vec<f32> = raw_coeffs.iter().map(|&v| v as f32 * scale_c).collect();
3017        let proto_data_f32 = ProtoData {
3018            mask_coefficients: make_f32(&[n_det, NUM_PROTOS], &coeffs_f32),
3019            protos: make_f32(&[PROTO_H, PROTO_W, NUM_PROTOS], &protos_f32),
3020            layout: ProtoLayout::Nhwc,
3021        };
3022
3023        let proto_data_int = ProtoData {
3024            mask_coefficients: make_i16_quant(&[n_det, NUM_PROTOS], &raw_coeffs, scale_c, 0),
3025            protos: make_i8_quant(&[PROTO_H, PROTO_W, NUM_PROTOS], &raw_protos, scale_p, 0),
3026            layout: ProtoLayout::Nhwc,
3027        };
3028
3029        let segs_f32 = cpu
3030            .materialize_segmentations(&detect, &proto_data_f32, None)
3031            .unwrap();
3032        let segs_int = cpu
3033            .materialize_segmentations(&detect, &proto_data_int, None)
3034            .unwrap();
3035
3036        assert_eq!(segs_f32.len(), segs_int.len());
3037        for (sf, si) in segs_f32.iter().zip(segs_int.iter()) {
3038            assert_eq!(sf.segmentation.shape(), si.segmentation.shape());
3039            let total = sf.segmentation.len();
3040            let mismatches = sf
3041                .segmentation
3042                .iter()
3043                .zip(si.segmentation.iter())
3044                .filter(|(a, b)| a != b)
3045                .count();
3046            let pct = mismatches as f64 / total as f64 * 100.0;
3047            assert!(
3048                pct < 5.0,
3049                "mask mismatch {mismatches}/{total} ({pct:.1}%) exceeds 5% threshold"
3050            );
3051        }
3052    }
3053
3054    // ── Multiple detections ────────────────────────────────────────
3055
3056    #[test]
3057    fn materialize_proto_i16_multiple_detections() {
3058        let cpu = CPUProcessor::new();
3059        let detect = vec![
3060            det(0.0, 0.0, 0.5, 0.5),
3061            det(0.5, 0.5, 1.0, 1.0),
3062            det(0.1, 0.1, 0.3, 0.3),
3063        ];
3064        let protos = make_i8_quant(
3065            &[PROTO_H, PROTO_W, NUM_PROTOS],
3066            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3067            0.02,
3068            0,
3069        );
3070        let coeffs = make_i16_quant(&[3, NUM_PROTOS], &gen_coeffs_i16(3, NUM_PROTOS), 0.01, 0);
3071        let proto_data = ProtoData {
3072            mask_coefficients: coeffs,
3073            protos,
3074            layout: ProtoLayout::Nhwc,
3075        };
3076        let segs = cpu
3077            .materialize_segmentations(&detect, &proto_data, None)
3078            .unwrap();
3079        assert_eq!(segs.len(), 3);
3080    }
3081
3082    // ── Empty detections ───────────────────────────────────────────
3083
3084    #[test]
3085    fn materialize_proto_i16_empty_detections() {
3086        let cpu = CPUProcessor::new();
3087        let detect: Vec<DetectBox> = vec![];
3088        let protos = make_i8_quant(
3089            &[PROTO_H, PROTO_W, NUM_PROTOS],
3090            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3091            0.02,
3092            0,
3093        );
3094        let coeffs = make_i16_quant(&[0, NUM_PROTOS], &[], 0.01, 0);
3095        let proto_data = ProtoData {
3096            mask_coefficients: coeffs,
3097            protos,
3098            layout: ProtoLayout::Nhwc,
3099        };
3100        let segs = cpu
3101            .materialize_segmentations(&detect, &proto_data, None)
3102            .unwrap();
3103        assert!(segs.is_empty());
3104    }
3105
3106    // ── Scaled parity ──────────────────────────────────────────────
3107
3108    #[test]
3109    fn materialize_scaled_i16_i8_matches_f32_reference() {
3110        let cpu = CPUProcessor::new();
3111        let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3112        let scale_c = 0.01_f32;
3113        let scale_p = 0.02_f32;
3114        let raw_protos = gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS);
3115        let raw_coeffs = gen_coeffs_i16(1, NUM_PROTOS);
3116
3117        let protos_f32: Vec<f32> = raw_protos.iter().map(|&v| v as f32 * scale_p).collect();
3118        let coeffs_f32: Vec<f32> = raw_coeffs.iter().map(|&v| v as f32 * scale_c).collect();
3119        let proto_data_f32 = ProtoData {
3120            mask_coefficients: make_f32(&[1, NUM_PROTOS], &coeffs_f32),
3121            protos: make_f32(&[PROTO_H, PROTO_W, NUM_PROTOS], &protos_f32),
3122            layout: ProtoLayout::Nhwc,
3123        };
3124        let proto_data_int = ProtoData {
3125            mask_coefficients: make_i16_quant(&[1, NUM_PROTOS], &raw_coeffs, scale_c, 0),
3126            protos: make_i8_quant(&[PROTO_H, PROTO_W, NUM_PROTOS], &raw_protos, scale_p, 0),
3127            layout: ProtoLayout::Nhwc,
3128        };
3129
3130        let (w, h) = (64_u32, 64_u32);
3131        let segs_f32 = cpu
3132            .materialize_scaled_segmentations(&detect, &proto_data_f32, None, w, h)
3133            .unwrap();
3134        let segs_int = cpu
3135            .materialize_scaled_segmentations(&detect, &proto_data_int, None, w, h)
3136            .unwrap();
3137
3138        assert_eq!(segs_f32.len(), segs_int.len());
3139        for (sf, si) in segs_f32.iter().zip(segs_int.iter()) {
3140            assert_eq!(sf.segmentation.shape(), si.segmentation.shape());
3141            let total = sf.segmentation.len();
3142            let mismatches = sf
3143                .segmentation
3144                .iter()
3145                .zip(si.segmentation.iter())
3146                .filter(|(a, b)| a != b)
3147                .count();
3148            let pct = mismatches as f64 / total as f64 * 100.0;
3149            assert!(
3150                pct < 5.0,
3151                "scaled mask mismatch {mismatches}/{total} ({pct:.1}%) exceeds 5% threshold"
3152            );
3153        }
3154    }
3155
3156    // ── i8×i8 existing path still works (regression) ───────────────
3157
3158    #[test]
3159    fn materialize_proto_i8_i8_regression() {
3160        let cpu = CPUProcessor::new();
3161        let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3162        let protos = make_i8_quant(
3163            &[PROTO_H, PROTO_W, NUM_PROTOS],
3164            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3165            0.02,
3166            0,
3167        );
3168        let coeffs = make_i8_quant(&[1, NUM_PROTOS], &gen_coeffs_i8(1, NUM_PROTOS), 0.01, 0);
3169        let proto_data = ProtoData {
3170            mask_coefficients: coeffs,
3171            protos,
3172            layout: ProtoLayout::Nhwc,
3173        };
3174        let result = cpu.materialize_segmentations(&detect, &proto_data, None);
3175        assert!(result.is_ok(), "i8×i8 regression: {:?}", result.err());
3176        assert_eq!(result.unwrap().len(), 1);
3177    }
3178
3179    // ── Non-zero zero_point ────────────────────────────────────────
3180
3181    #[test]
3182    fn materialize_proto_i16_nonzero_zp() {
3183        let cpu = CPUProcessor::new();
3184        let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3185        let protos = make_i8_quant(
3186            &[PROTO_H, PROTO_W, NUM_PROTOS],
3187            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3188            0.02,
3189            -10,
3190        );
3191        let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 5);
3192        let proto_data = ProtoData {
3193            mask_coefficients: coeffs,
3194            protos,
3195            layout: ProtoLayout::Nhwc,
3196        };
3197        let result = cpu.materialize_segmentations(&detect, &proto_data, None);
3198        assert!(result.is_ok(), "nonzero zp failed: {:?}", result.err());
3199        assert_eq!(result.unwrap().len(), 1);
3200    }
3201
3202    // ── Scaled: non-zero zero_point ────────────────────────────────
3203
3204    #[test]
3205    fn materialize_scaled_i16_nonzero_zp() {
3206        let cpu = CPUProcessor::new();
3207        let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3208        let protos = make_i8_quant(
3209            &[PROTO_H, PROTO_W, NUM_PROTOS],
3210            &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3211            0.02,
3212            -10,
3213        );
3214        let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 5);
3215        let proto_data = ProtoData {
3216            mask_coefficients: coeffs,
3217            protos,
3218            layout: ProtoLayout::Nhwc,
3219        };
3220        let result = cpu.materialize_scaled_segmentations(&detect, &proto_data, None, 64, 64);
3221        assert!(
3222            result.is_ok(),
3223            "scaled nonzero zp failed: {:?}",
3224            result.err()
3225        );
3226        assert_eq!(result.unwrap().len(), 1);
3227    }
3228}