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