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