1use 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 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 if seg_classes > u16::MAX as usize {
51 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 return Err(crate::Error::InvalidShape(format!(
129 "segmentation must be [H, W, C], got {:?}",
130 seg.shape()
131 )));
132 };
133 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 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 #[allow(clippy::too_many_arguments)]
211 pub(super) fn render_box(
212 &mut self,
213 dst_w: usize,
214 dst_h: usize,
215 dst_rs: usize,
216 dst_c: usize,
217 dst_slice: &mut [u8],
218 detect: &[DetectBox],
219 color_mode: crate::ColorMode,
220 ) -> Result<()> {
221 const LINE_THICKNESS: usize = 3;
222
223 for (idx, d) in detect.iter().enumerate() {
224 use edgefirst_tensor::BoundingBox;
225
226 let color_index = color_mode.index(idx, d.label);
227 let [r, g, b, _] = self.colors[color_index % self.colors.len()];
228 let bbox = d.bbox.to_canonical();
229 let bbox = BoundingBox {
230 xmin: bbox.xmin.clamp(0.0, 1.0),
231 ymin: bbox.ymin.clamp(0.0, 1.0),
232 xmax: bbox.xmax.clamp(0.0, 1.0),
233 ymax: bbox.ymax.clamp(0.0, 1.0),
234 };
235 let inner = [
236 ((dst_w - 1) as f32 * bbox.xmin - 0.5).round() as usize,
237 ((dst_h - 1) as f32 * bbox.ymin - 0.5).round() as usize,
238 ((dst_w - 1) as f32 * bbox.xmax + 0.5).round() as usize,
239 ((dst_h - 1) as f32 * bbox.ymax + 0.5).round() as usize,
240 ];
241
242 let outer = [
243 inner[0].saturating_sub(LINE_THICKNESS),
244 inner[1].saturating_sub(LINE_THICKNESS),
245 (inner[2] + LINE_THICKNESS).min(dst_w),
246 (inner[3] + LINE_THICKNESS).min(dst_h),
247 ];
248
249 for y in outer[1] + 1..=inner[1] {
251 for x in outer[0] + 1..outer[2] {
252 let index = (y * dst_rs) + (x * dst_c);
253 dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
254 }
255 }
256
257 for y in inner[1]..inner[3] {
259 for x in outer[0] + 1..=inner[0] {
260 let index = (y * dst_rs) + (x * dst_c);
261 dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
262 }
263
264 for x in inner[2]..outer[2] {
265 let index = (y * dst_rs) + (x * dst_c);
266 dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
267 }
268 }
269
270 for y in inner[3]..outer[3] {
272 for x in outer[0] + 1..outer[2] {
273 let index = (y * dst_rs) + (x * dst_c);
274 dst_slice[index..(index + 3)].copy_from_slice(&[r, g, b]);
275 }
276 }
277 }
278 Ok(())
279 }
280
281 pub fn materialize_segmentations(
292 &self,
293 detect: &[crate::DetectBox],
294 proto_data: &crate::ProtoData,
295 letterbox: Option<[f32; 4]>,
296 ) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
297 use edgefirst_tensor::{DType, TensorMapTrait, TensorTrait};
298
299 let _span = tracing::trace_span!(
300 "image.materialize_masks",
301 mode = "proto",
302 n_detections = detect.len(),
303 )
304 .entered();
305
306 if detect.is_empty() {
307 return Ok(Vec::new());
308 }
309 let proto_shape = proto_data.protos.shape();
310 if proto_shape.len() != 3 {
311 return Err(crate::Error::InvalidShape(format!(
312 "protos tensor must be rank-3, got {proto_shape:?}"
313 )));
314 }
315 let (proto_h, proto_w, num_protos) = match proto_data.layout {
317 edgefirst_tensor::ProtoLayout::Nhwc => (proto_shape[0], proto_shape[1], proto_shape[2]),
318 edgefirst_tensor::ProtoLayout::Nchw => (proto_shape[1], proto_shape[2], proto_shape[0]),
319 };
320 let coeff_shape = proto_data.mask_coefficients.shape();
321 if coeff_shape.len() != 2 || coeff_shape[1] != num_protos {
322 return Err(crate::Error::InvalidShape(format!(
323 "mask_coefficients shape {coeff_shape:?} incompatible with protos \
324 {proto_shape:?} (expected [N, {num_protos}])"
325 )));
326 }
327 if coeff_shape[0] == 0 {
328 return Ok(Vec::new());
329 }
330 if coeff_shape[0] != detect.len() {
331 return Err(crate::Error::Internal(format!(
332 "mask_coefficients rows {} != detection count {}",
333 coeff_shape[0],
334 detect.len()
335 )));
336 }
337
338 let (lx0, inv_lw, ly0, inv_lh) = match letterbox {
340 Some([lx0, ly0, lx1, ly1]) => {
341 let lw = lx1 - lx0;
342 let lh = ly1 - ly0;
343 (
344 lx0,
345 if lw > 0.0 { 1.0 / lw } else { 1.0 },
346 ly0,
347 if lh > 0.0 { 1.0 / lh } else { 1.0 },
348 )
349 }
350 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
351 };
352
353 if proto_data.mask_coefficients.dtype() == DType::I8
360 && proto_data.protos.dtype() == DType::I8
361 {
362 let coeff_t = proto_data
363 .mask_coefficients
364 .as_typed::<i8>()
365 .expect("I8 coefficients");
366 let coeff_m = coeff_t.map_read()?;
367 let coeff_quant = coeff_t.quantization().ok_or_else(|| {
368 crate::Error::InvalidShape(
369 "I8 mask_coefficients require quantization metadata".into(),
370 )
371 })?;
372 let proto_t = proto_data.protos.as_typed::<i8>().expect("I8 protos");
373 let proto_m = proto_t.map_read()?;
374 let proto_quant = proto_t.quantization().ok_or_else(|| {
375 crate::Error::InvalidShape("I8 protos require quantization metadata".into())
376 })?;
377 match proto_segmentations_i8_i8(
378 detect,
379 coeff_m.as_slice(),
380 coeff_quant,
381 proto_m.as_slice(),
382 proto_quant,
383 proto_h,
384 proto_w,
385 num_protos,
386 lx0,
387 inv_lw,
388 ly0,
389 inv_lh,
390 proto_data.layout,
391 ) {
392 Ok(result) => return Ok(result),
393 Err(crate::Error::NotSupported(_)) => {
394 }
397 Err(e) => return Err(e),
398 }
399 }
400
401 if proto_data.mask_coefficients.dtype() == DType::I16
403 && proto_data.protos.dtype() == DType::I8
404 {
405 let coeff_t = proto_data
406 .mask_coefficients
407 .as_typed::<i16>()
408 .expect("I16 coefficients");
409 let coeff_m = coeff_t.map_read()?;
410 if let Some(coeff_quant) = coeff_t.quantization() {
413 let proto_t = proto_data.protos.as_typed::<i8>().expect("I8 protos");
414 let proto_m = proto_t.map_read()?;
415 let proto_quant = proto_t.quantization().ok_or_else(|| {
416 crate::Error::InvalidShape("I8 protos require quantization metadata".into())
417 })?;
418 match proto_segmentations_i16_i8(
419 detect,
420 coeff_m.as_slice(),
421 coeff_quant,
422 proto_m.as_slice(),
423 proto_quant,
424 proto_h,
425 proto_w,
426 num_protos,
427 lx0,
428 inv_lw,
429 ly0,
430 inv_lh,
431 proto_data.layout,
432 ) {
433 Ok(result) => return Ok(result),
434 Err(crate::Error::NotSupported(_)) => {
435 }
437 Err(e) => return Err(e),
438 }
439 }
440 }
441
442 if proto_data.layout == edgefirst_tensor::ProtoLayout::Nchw
447 && proto_data.protos.dtype() != DType::I8
448 {
449 return Err(crate::Error::NotSupported(
450 "NCHW proto layout with non-I8 protos is not supported in the f32 fallback path"
451 .into(),
452 ));
453 }
454 let coeff_f32_storage: Vec<f32>;
455 let coeff_f32_slice: &[f32] = match proto_data.mask_coefficients.dtype() {
456 DType::F32 => {
457 let t = proto_data
458 .mask_coefficients
459 .as_typed::<f32>()
460 .expect("dtype matched F32");
461 let m = t.map_read()?;
462 coeff_f32_storage = m.as_slice().to_vec();
463 &coeff_f32_storage[..]
464 }
465 DType::F16 => {
466 let t = proto_data
467 .mask_coefficients
468 .as_typed::<half::f16>()
469 .expect("dtype matched F16");
470 let m = t.map_read()?;
471 coeff_f32_storage = m.as_slice().iter().map(|v| v.to_f32()).collect();
472 &coeff_f32_storage[..]
473 }
474 DType::I8 => {
475 let t = proto_data
476 .mask_coefficients
477 .as_typed::<i8>()
478 .expect("dtype matched I8");
479 let m = t.map_read()?;
480 coeff_f32_storage = if let Some(q) = t.quantization() {
481 use edgefirst_tensor::QuantMode;
482 let (scale, zp) = match q.mode() {
483 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
484 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
485 other => {
486 return Err(crate::Error::NotSupported(format!(
487 "I8 mask_coefficients quantization mode {other:?} not supported"
488 )));
489 }
490 };
491 m.as_slice()
492 .iter()
493 .map(|&v| (v as f32 - zp) * scale)
494 .collect()
495 } else {
496 m.as_slice().iter().map(|&v| v as f32).collect()
497 };
498 &coeff_f32_storage[..]
499 }
500 DType::I16 => {
501 let t = proto_data
502 .mask_coefficients
503 .as_typed::<i16>()
504 .expect("dtype matched I16");
505 let m = t.map_read()?;
506 coeff_f32_storage = if let Some(q) = t.quantization() {
507 use edgefirst_tensor::QuantMode;
508 let (scale, zp) = match q.mode() {
509 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
510 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
511 other => {
512 return Err(crate::Error::NotSupported(format!(
513 "I16 mask_coefficients quantization mode {other:?} not supported"
514 )));
515 }
516 };
517 m.as_slice()
518 .iter()
519 .map(|&v| (v as f32 - zp) * scale)
520 .collect()
521 } else {
522 m.as_slice().iter().map(|&v| v as f32).collect()
523 };
524 &coeff_f32_storage[..]
525 }
526 other => {
527 return Err(crate::Error::InvalidShape(format!(
528 "mask_coefficients dtype {other:?} not supported; expected F32, F16, I8, or I16"
529 )));
530 }
531 };
532
533 match proto_data.protos.dtype() {
539 DType::I8 => {
540 let t = proto_data
541 .protos
542 .as_typed::<i8>()
543 .expect("dtype matched I8");
544 let quant = t.quantization().ok_or_else(|| {
545 crate::Error::InvalidShape("I8 protos require quantization metadata".into())
546 })?;
547 let m = t.map_read()?;
548 let src_slice = m.as_slice();
549 let transposed_storage = if proto_data.layout == edgefirst_tensor::ProtoLayout::Nchw
550 {
551 let hw = proto_h.checked_mul(proto_w).ok_or_else(|| {
556 crate::Error::InvalidShape(format!(
557 "proto plane size overflow (proto_h={proto_h}, proto_w={proto_w})"
558 ))
559 })?;
560 let need = hw.checked_mul(num_protos).ok_or_else(|| {
561 crate::Error::InvalidShape(format!(
562 "proto NCHW size overflow (hw={hw}, n={num_protos})"
563 ))
564 })?;
565 if src_slice.len() < need {
566 return Err(crate::Error::InvalidShape(format!(
567 "proto buffer {} bytes < {need} (proto_h={proto_h}, \
568 proto_w={proto_w}, num_protos={num_protos})",
569 src_slice.len()
570 )));
571 }
572 let mut nhwc = vec![0i8; need];
573 for c in 0..num_protos {
574 let plane = &src_slice[c * hw..(c + 1) * hw];
575 for px in 0..hw {
576 nhwc[px * num_protos + c] = plane[px];
577 }
578 }
579 Some(nhwc)
580 } else {
581 None
582 };
583 let protos_slice = transposed_storage.as_deref().unwrap_or(src_slice);
584 detect
585 .par_iter()
586 .enumerate()
587 .map(|(i, det)| {
588 let coeff = &coeff_f32_slice[i * num_protos..(i + 1) * num_protos];
589 let (x0, y0, x1, y1, roi_w, roi_h) =
590 bbox_to_proto_roi(det, proto_w, proto_h);
591 let mask = fused_dequant_dot_sign_i8_slice(
592 protos_slice,
593 coeff,
594 quant,
595 proto_h,
596 proto_w,
597 y0,
598 x0,
599 roi_h,
600 roi_w,
601 num_protos,
602 )?;
603 seg_from_roi(
604 mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
605 )
606 })
607 .collect()
608 }
609 DType::F32 => {
610 let t = proto_data
611 .protos
612 .as_typed::<f32>()
613 .expect("dtype matched F32");
614 let m = t.map_read()?;
615 let protos_slice = m.as_slice();
616 detect
617 .par_iter()
618 .enumerate()
619 .map(|(i, det)| {
620 let coeff = &coeff_f32_slice[i * num_protos..(i + 1) * num_protos];
621 let (x0, y0, x1, y1, roi_w, roi_h) =
622 bbox_to_proto_roi(det, proto_w, proto_h);
623 let mask = fused_dot_sign_f32_slice(
624 protos_slice,
625 coeff,
626 proto_h,
627 proto_w,
628 y0,
629 x0,
630 roi_h,
631 roi_w,
632 num_protos,
633 );
634 seg_from_roi(
635 mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
636 )
637 })
638 .collect()
639 }
640 DType::F16 => {
641 let t = proto_data
642 .protos
643 .as_typed::<half::f16>()
644 .expect("dtype matched F16");
645 let m = t.map_read()?;
646 let protos_slice = m.as_slice();
647 detect
648 .par_iter()
649 .enumerate()
650 .map(|(i, det)| {
651 let coeff = &coeff_f32_slice[i * num_protos..(i + 1) * num_protos];
652 let (x0, y0, x1, y1, roi_w, roi_h) =
653 bbox_to_proto_roi(det, proto_w, proto_h);
654 let mask = fused_dot_sign_f16_slice(
655 protos_slice,
656 coeff,
657 proto_h,
658 proto_w,
659 y0,
660 x0,
661 roi_h,
662 roi_w,
663 num_protos,
664 );
665 seg_from_roi(
666 mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
667 )
668 })
669 .collect()
670 }
671 other => Err(crate::Error::InvalidShape(format!(
672 "proto tensor dtype {other:?} not supported"
673 ))),
674 }
675 }
676
677 pub fn materialize_scaled_segmentations(
690 &self,
691 detect: &[crate::DetectBox],
692 proto_data: &crate::ProtoData,
693 letterbox: Option<[f32; 4]>,
694 width: u32,
695 height: u32,
696 ) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
697 use edgefirst_tensor::{DType, TensorMapTrait, TensorTrait};
698
699 let _span = tracing::trace_span!(
700 "image.materialize_masks",
701 mode = "scaled",
702 n_detections = detect.len(),
703 width,
704 height,
705 )
706 .entered();
707
708 if detect.is_empty() {
709 return Ok(Vec::new());
710 }
711 if width == 0 || height == 0 {
712 return Err(crate::Error::InvalidShape(
713 "Scaled mask width/height must be positive".into(),
714 ));
715 }
716 let proto_shape = proto_data.protos.shape();
717 if proto_shape.len() != 3 {
718 return Err(crate::Error::InvalidShape(format!(
719 "protos tensor must be rank-3, got {proto_shape:?}"
720 )));
721 }
722 let (proto_h, proto_w, num_protos) = match proto_data.layout {
724 edgefirst_tensor::ProtoLayout::Nhwc => (proto_shape[0], proto_shape[1], proto_shape[2]),
725 edgefirst_tensor::ProtoLayout::Nchw => (proto_shape[1], proto_shape[2], proto_shape[0]),
726 };
727 let coeff_shape = proto_data.mask_coefficients.shape();
728 if coeff_shape.len() != 2 || coeff_shape[1] != num_protos {
729 return Err(crate::Error::InvalidShape(format!(
730 "mask_coefficients shape {coeff_shape:?} incompatible with protos \
731 {proto_shape:?}"
732 )));
733 }
734 if coeff_shape[0] == 0 {
735 return Ok(Vec::new());
736 }
737 if coeff_shape[0] != detect.len() {
738 return Err(crate::Error::Internal(format!(
739 "mask_coefficients rows {} != detection count {}",
740 coeff_shape[0],
741 detect.len()
742 )));
743 }
744
745 if proto_data.mask_coefficients.dtype() == DType::I8
751 && proto_data.protos.dtype() == DType::I8
752 {
753 let coeff_t = proto_data
754 .mask_coefficients
755 .as_typed::<i8>()
756 .expect("I8 coefficients");
757 let coeff_m = coeff_t.map_read()?;
758 let coeff_quant = coeff_t.quantization().ok_or_else(|| {
759 crate::Error::InvalidShape(
760 "I8 mask_coefficients require quantization metadata".into(),
761 )
762 })?;
763 let proto_t = proto_data.protos.as_typed::<i8>().expect("I8 protos");
764 let proto_m = proto_t.map_read()?;
765 let proto_quant = proto_t.quantization().ok_or_else(|| {
766 crate::Error::InvalidShape("I8 protos require quantization metadata".into())
767 })?;
768 match scaled_segmentations_i8_i8(
769 detect,
770 coeff_m.as_slice(),
771 coeff_quant,
772 proto_m.as_slice(),
773 proto_quant,
774 proto_h,
775 proto_w,
776 num_protos,
777 letterbox,
778 width,
779 height,
780 proto_data.layout,
781 ) {
782 Ok(result) => return Ok(result),
783 Err(crate::Error::NotSupported(_)) => {
784 }
787 Err(e) => return Err(e),
788 }
789 }
790
791 if proto_data.mask_coefficients.dtype() == DType::I16
793 && proto_data.protos.dtype() == DType::I8
794 {
795 let coeff_t = proto_data
796 .mask_coefficients
797 .as_typed::<i16>()
798 .expect("I16 coefficients");
799 let coeff_m = coeff_t.map_read()?;
800 if let Some(coeff_quant) = coeff_t.quantization() {
803 let proto_t = proto_data.protos.as_typed::<i8>().expect("I8 protos");
804 let proto_m = proto_t.map_read()?;
805 let proto_quant = proto_t.quantization().ok_or_else(|| {
806 crate::Error::InvalidShape("I8 protos require quantization metadata".into())
807 })?;
808 match scaled_segmentations_i16_i8(
809 detect,
810 coeff_m.as_slice(),
811 coeff_quant,
812 proto_m.as_slice(),
813 proto_quant,
814 proto_h,
815 proto_w,
816 num_protos,
817 letterbox,
818 width,
819 height,
820 proto_data.layout,
821 ) {
822 Ok(result) => return Ok(result),
823 Err(crate::Error::NotSupported(_)) => {}
824 Err(e) => return Err(e),
825 }
826 }
827 }
828
829 if proto_data.layout == edgefirst_tensor::ProtoLayout::Nchw
831 && proto_data.protos.dtype() != DType::I8
832 {
833 return Err(crate::Error::NotSupported(
834 "NCHW proto layout with non-I8 protos is not supported in the f32 fallback path"
835 .into(),
836 ));
837 }
838 let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
839 DType::F32 => {
840 let t = proto_data.mask_coefficients.as_typed::<f32>().expect("F32");
841 let m = t.map_read()?;
842 m.as_slice().to_vec()
843 }
844 DType::F16 => {
845 let t = proto_data
846 .mask_coefficients
847 .as_typed::<half::f16>()
848 .expect("F16");
849 let m = t.map_read()?;
850 m.as_slice().iter().map(|v| v.to_f32()).collect()
851 }
852 DType::I8 => {
853 let t = proto_data.mask_coefficients.as_typed::<i8>().expect("I8");
855 let m = t.map_read()?;
856 let q = t.quantization().ok_or_else(|| {
857 crate::Error::InvalidShape(
858 "I8 mask_coefficients require quantization metadata".into(),
859 )
860 })?;
861 use edgefirst_tensor::QuantMode;
862 let (scale, zp) = match q.mode() {
863 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
864 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
865 _ => {
866 return Err(crate::Error::NotSupported(
867 "per-channel mask_coefficients not supported".into(),
868 ))
869 }
870 };
871 m.as_slice()
872 .iter()
873 .map(|&v| (v as f32 - zp) * scale)
874 .collect()
875 }
876 DType::I16 => {
877 let t = proto_data.mask_coefficients.as_typed::<i16>().expect("I16");
878 let m = t.map_read()?;
879 if let Some(q) = t.quantization() {
880 use edgefirst_tensor::QuantMode;
881 let (scale, zp) = match q.mode() {
882 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
883 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
884 other => {
885 return Err(crate::Error::NotSupported(format!(
886 "I16 mask_coefficients quantization mode {other:?} not supported"
887 )))
888 }
889 };
890 m.as_slice()
891 .iter()
892 .map(|&v| (v as f32 - zp) * scale)
893 .collect()
894 } else {
895 m.as_slice().iter().map(|&v| v as f32).collect()
896 }
897 }
898 other => {
899 return Err(crate::Error::InvalidShape(format!(
900 "mask_coefficients dtype {other:?} not supported"
901 )));
902 }
903 };
904
905 match proto_data.protos.dtype() {
906 DType::F32 => {
907 let t = proto_data.protos.as_typed::<f32>().expect("F32");
908 let m = t.map_read()?;
909 scaled_segmentations_f32_slice(
910 detect,
911 &coeff_f32,
912 m.as_slice(),
913 proto_h,
914 proto_w,
915 num_protos,
916 letterbox,
917 width,
918 height,
919 )
920 }
921 DType::F16 => {
922 let t = proto_data.protos.as_typed::<half::f16>().expect("F16");
923 let m = t.map_read()?;
924 scaled_segmentations_f16_slice(
925 detect,
926 &coeff_f32,
927 m.as_slice(),
928 proto_h,
929 proto_w,
930 num_protos,
931 letterbox,
932 width,
933 height,
934 )
935 }
936 DType::I8 => {
937 let t = proto_data.protos.as_typed::<i8>().expect("I8");
938 let m = t.map_read()?;
939 let quant = t.quantization().ok_or_else(|| {
940 crate::Error::InvalidShape("I8 protos require quantization metadata".into())
941 })?;
942 let src_slice = m.as_slice();
943 let transposed_storage = if proto_data.layout == edgefirst_tensor::ProtoLayout::Nchw
944 {
945 let hw = proto_h.checked_mul(proto_w).ok_or_else(|| {
950 crate::Error::InvalidShape(format!(
951 "proto plane size overflow (proto_h={proto_h}, proto_w={proto_w})"
952 ))
953 })?;
954 let need = hw.checked_mul(num_protos).ok_or_else(|| {
955 crate::Error::InvalidShape(format!(
956 "proto NCHW size overflow (hw={hw}, n={num_protos})"
957 ))
958 })?;
959 if src_slice.len() < need {
960 return Err(crate::Error::InvalidShape(format!(
961 "proto buffer {} bytes < {need} (proto_h={proto_h}, \
962 proto_w={proto_w}, num_protos={num_protos})",
963 src_slice.len()
964 )));
965 }
966 let mut nhwc = vec![0i8; need];
967 for c in 0..num_protos {
968 let plane = &src_slice[c * hw..(c + 1) * hw];
969 for px in 0..hw {
970 nhwc[px * num_protos + c] = plane[px];
971 }
972 }
973 Some(nhwc)
974 } else {
975 None
976 };
977 let protos_slice = transposed_storage.as_deref().unwrap_or(src_slice);
978 scaled_segmentations_i8_slice(
979 detect,
980 &coeff_f32,
981 protos_slice,
982 proto_h,
983 proto_w,
984 num_protos,
985 quant,
986 letterbox,
987 width,
988 height,
989 )
990 }
991 other => Err(crate::Error::InvalidShape(format!(
992 "proto tensor dtype {other:?} not supported"
993 ))),
994 }
995 }
996}
997
998fn mask_to_tensor(mask: ndarray::ArrayView3<'_, u8>) -> Result<edgefirst_tensor::TensorDyn> {
1022 edgefirst_tensor::Tensor::from_arrayview3(mask)
1023 .map(edgefirst_tensor::TensorDyn::from)
1024 .map_err(|e| crate::Error::InvalidShape(format!("mask -> TensorDyn: {e}")))
1025}
1026
1027fn bbox_to_proto_roi(
1028 det: &DetectBox,
1029 proto_w: usize,
1030 proto_h: usize,
1031) -> (usize, usize, usize, usize, usize, usize) {
1032 let bbox = det.bbox.to_canonical();
1033 let xmin = bbox.xmin.clamp(0.0, 1.0);
1034 let ymin = bbox.ymin.clamp(0.0, 1.0);
1035 let xmax = bbox.xmax.clamp(0.0, 1.0);
1036 let ymax = bbox.ymax.clamp(0.0, 1.0);
1037 let x0 = ((xmin * proto_w as f32) as usize).min(proto_w.saturating_sub(1));
1038 let y0 = ((ymin * proto_h as f32) as usize).min(proto_h.saturating_sub(1));
1039 let x1 = ((xmax * proto_w as f32).ceil() as usize).min(proto_w);
1040 let y1 = ((ymax * proto_h as f32).ceil() as usize).min(proto_h);
1041 let roi_w = x1.saturating_sub(x0).max(1);
1042 let roi_h = y1.saturating_sub(y0).max(1);
1043 (x0, y0, x1, y1, roi_w, roi_h)
1044}
1045
1046#[allow(clippy::too_many_arguments)]
1050fn seg_from_roi(
1051 mask: ndarray::Array3<u8>,
1052 x0: usize,
1053 y0: usize,
1054 x1: usize,
1055 y1: usize,
1056 proto_w: usize,
1057 proto_h: usize,
1058 lx0: f32,
1059 inv_lw: f32,
1060 ly0: f32,
1061 inv_lh: f32,
1062) -> Result<Segmentation> {
1063 let seg_xmin = ((x0 as f32 / proto_w as f32) - lx0) * inv_lw;
1064 let seg_ymin = ((y0 as f32 / proto_h as f32) - ly0) * inv_lh;
1065 let seg_xmax = ((x1 as f32 / proto_w as f32) - lx0) * inv_lw;
1066 let seg_ymax = ((y1 as f32 / proto_h as f32) - ly0) * inv_lh;
1067 Ok(Segmentation {
1068 xmin: seg_xmin.clamp(0.0, 1.0),
1069 ymin: seg_ymin.clamp(0.0, 1.0),
1070 xmax: seg_xmax.clamp(0.0, 1.0),
1071 ymax: seg_ymax.clamp(0.0, 1.0),
1072 segmentation: mask_to_tensor(mask.view())?,
1073 })
1074}
1075
1076#[allow(clippy::too_many_arguments)]
1092fn proto_segmentations_i8_i8(
1093 detect: &[crate::DetectBox],
1094 coeff_all: &[i8],
1095 coeff_quant: &edgefirst_tensor::Quantization,
1096 protos: &[i8],
1097 proto_quant: &edgefirst_tensor::Quantization,
1098 proto_h: usize,
1099 proto_w: usize,
1100 num_protos: usize,
1101 lx0: f32,
1102 inv_lw: f32,
1103 ly0: f32,
1104 inv_lh: f32,
1105 layout: edgefirst_tensor::ProtoLayout,
1106) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
1107 use edgefirst_tensor::QuantMode;
1108
1109 let _span = tracing::trace_span!(
1110 "image.materialize_masks.kernel_i8",
1111 n = detect.len(),
1112 proto_h,
1113 proto_w,
1114 num_protos,
1115 ?layout,
1116 )
1117 .entered();
1118
1119 let zp_c: i32 = match coeff_quant.mode() {
1120 QuantMode::PerTensor { zero_point, .. } => zero_point,
1121 QuantMode::PerTensorSymmetric { .. } => 0,
1122 _ => {
1123 return Err(crate::Error::NotSupported(
1124 "per-channel coeff quantization not supported on proto-res i8 path".into(),
1125 ))
1126 }
1127 };
1128 let zp_p: i32 = match proto_quant.mode() {
1129 QuantMode::PerTensor { zero_point, .. } => zero_point,
1130 QuantMode::PerTensorSymmetric { .. } => 0,
1131 _ => {
1132 return Err(crate::Error::NotSupported(
1133 "per-channel proto quantization not supported on proto-res i8 path".into(),
1134 ))
1135 }
1136 };
1137
1138 let hw = proto_h * proto_w;
1139
1140 let proto_sums: Vec<i32> = if zp_c != 0 {
1142 match layout {
1143 edgefirst_tensor::ProtoLayout::Nhwc => (0..hw)
1144 .map(|px_idx| {
1145 let base = px_idx * num_protos;
1146 protos[base..base + num_protos]
1147 .iter()
1148 .map(|&v| v as i32)
1149 .sum()
1150 })
1151 .collect(),
1152 edgefirst_tensor::ProtoLayout::Nchw => {
1153 let mut sums = vec![0i32; hw];
1154 for c in 0..num_protos {
1155 let plane = &protos[c * hw..];
1156 for (px, s) in sums.iter_mut().enumerate() {
1157 *s += plane[px] as i32;
1158 }
1159 }
1160 sums
1161 }
1162 }
1163 } else {
1164 Vec::new()
1165 };
1166
1167 #[cfg(target_arch = "aarch64")]
1168 let use_dotprod = std::arch::is_aarch64_feature_detected!("dotprod");
1169
1170 detect
1171 .par_iter()
1172 .enumerate()
1173 .map(|(i, det)| {
1174 let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
1175 let (x0, y0, x1, y1, roi_w, roi_h) = bbox_to_proto_roi(det, proto_w, proto_h);
1176
1177 let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
1179 let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
1180
1181 let mut mask_buf = vec![0u8; roi_h * roi_w];
1182
1183 match layout {
1184 edgefirst_tensor::ProtoLayout::Nhwc => {
1185 let stride_y = proto_w * num_protos;
1186 #[cfg(target_arch = "aarch64")]
1187 {
1188 if use_dotprod {
1189 for ly in 0..roi_h {
1190 let py = y0 + ly;
1191 let row_base = py * stride_y + x0 * num_protos;
1192 for lx in 0..roi_w {
1193 let pix_base = row_base + lx * num_protos;
1194 let proto_px = &protos[pix_base..pix_base + num_protos];
1195 let raw_dot = unsafe {
1196 dot_i8_neon_dotprod(
1197 coeff.as_ptr(),
1198 proto_px.as_ptr(),
1199 num_protos,
1200 )
1201 };
1202 let correction = if zp_c != 0 {
1203 zp_c * proto_sums[py * proto_w + x0 + lx]
1204 } else {
1205 0
1206 };
1207 let logit = raw_dot - correction - bias;
1208 if logit > 0 {
1209 mask_buf[ly * roi_w + lx] = 255;
1210 }
1211 }
1212 }
1213 } else {
1214 for ly in 0..roi_h {
1215 let py = y0 + ly;
1216 let row_base = py * stride_y + x0 * num_protos;
1217 for lx in 0..roi_w {
1218 let pix_base = row_base + lx * num_protos;
1219 let proto_px = &protos[pix_base..pix_base + num_protos];
1220 let raw_dot = unsafe {
1221 dot_i8_neon_base(
1222 coeff.as_ptr(),
1223 proto_px.as_ptr(),
1224 num_protos,
1225 )
1226 };
1227 let correction = if zp_c != 0 {
1228 zp_c * proto_sums[py * proto_w + x0 + lx]
1229 } else {
1230 0
1231 };
1232 let logit = raw_dot - correction - bias;
1233 if logit > 0 {
1234 mask_buf[ly * roi_w + lx] = 255;
1235 }
1236 }
1237 }
1238 }
1239 }
1240 #[cfg(not(target_arch = "aarch64"))]
1241 {
1242 for ly in 0..roi_h {
1243 let py = y0 + ly;
1244 let row_base = py * stride_y + x0 * num_protos;
1245 for lx in 0..roi_w {
1246 let pix_base = row_base + lx * num_protos;
1247 let proto_px = &protos[pix_base..pix_base + num_protos];
1248 let raw_dot = dot_i8_scalar(coeff, proto_px, num_protos);
1249 let correction = if zp_c != 0 {
1250 zp_c * proto_sums[py * proto_w + x0 + lx]
1251 } else {
1252 0
1253 };
1254 let logit = raw_dot - correction - bias;
1255 if logit > 0 {
1256 mask_buf[ly * roi_w + lx] = 255;
1257 }
1258 }
1259 }
1260 }
1261 }
1262 edgefirst_tensor::ProtoLayout::Nchw => {
1263 let mut accum = vec![0i32; roi_h * roi_w];
1267 for c in 0..num_protos {
1268 let plane = &protos[c * hw..];
1269 let coeff_c = coeff[c] as i32;
1270 for ly in 0..roi_h {
1271 let py = y0 + ly;
1272 let row_start = py * proto_w + x0;
1273 let out_row_start = ly * roi_w;
1274 for lx in 0..roi_w {
1275 accum[out_row_start + lx] += coeff_c * plane[row_start + lx] as i32;
1276 }
1277 }
1278 }
1279 for ly in 0..roi_h {
1281 let py = y0 + ly;
1282 for lx in 0..roi_w {
1283 let idx = ly * roi_w + lx;
1284 let correction = if zp_c != 0 {
1285 zp_c * proto_sums[py * proto_w + x0 + lx]
1286 } else {
1287 0
1288 };
1289 let logit = accum[idx] - correction - bias;
1290 if logit > 0 {
1291 mask_buf[idx] = 255;
1292 }
1293 }
1294 }
1295 }
1296 }
1297
1298 let mask = ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1299 .expect("mask_buf length matches roi_h * roi_w");
1300 seg_from_roi(
1301 mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
1302 )
1303 })
1304 .collect()
1305}
1306
1307#[allow(clippy::too_many_arguments)]
1309fn proto_segmentations_i16_i8(
1310 detect: &[crate::DetectBox],
1311 coeff_all: &[i16],
1312 coeff_quant: &edgefirst_tensor::Quantization,
1313 protos: &[i8],
1314 proto_quant: &edgefirst_tensor::Quantization,
1315 proto_h: usize,
1316 proto_w: usize,
1317 num_protos: usize,
1318 lx0: f32,
1319 inv_lw: f32,
1320 ly0: f32,
1321 inv_lh: f32,
1322 layout: edgefirst_tensor::ProtoLayout,
1323) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
1324 use edgefirst_tensor::QuantMode;
1325
1326 let _span = tracing::trace_span!(
1327 "image.materialize_masks.kernel_i16xi8",
1328 n = detect.len(),
1329 proto_h,
1330 proto_w,
1331 num_protos,
1332 ?layout,
1333 )
1334 .entered();
1335
1336 let zp_c: i32 = match coeff_quant.mode() {
1337 QuantMode::PerTensor { zero_point, .. } => zero_point,
1338 QuantMode::PerTensorSymmetric { .. } => 0,
1339 _ => {
1340 return Err(crate::Error::NotSupported(
1341 "per-channel coeff quantization not supported on proto-res i16 path".into(),
1342 ))
1343 }
1344 };
1345 let zp_p: i32 = match proto_quant.mode() {
1346 QuantMode::PerTensor { zero_point, .. } => zero_point,
1347 QuantMode::PerTensorSymmetric { .. } => 0,
1348 _ => {
1349 return Err(crate::Error::NotSupported(
1350 "per-channel proto quantization not supported on proto-res i8 path".into(),
1351 ))
1352 }
1353 };
1354
1355 let hw = proto_h * proto_w;
1356
1357 let proto_sums: Vec<i32> = if zp_c != 0 {
1359 match layout {
1360 edgefirst_tensor::ProtoLayout::Nhwc => (0..hw)
1361 .map(|px_idx| {
1362 let base = px_idx * num_protos;
1363 protos[base..base + num_protos]
1364 .iter()
1365 .map(|&v| v as i32)
1366 .sum()
1367 })
1368 .collect(),
1369 edgefirst_tensor::ProtoLayout::Nchw => {
1370 let mut sums = vec![0i32; hw];
1371 for c in 0..num_protos {
1372 let plane = &protos[c * hw..];
1373 for (px, s) in sums.iter_mut().enumerate() {
1374 *s += plane[px] as i32;
1375 }
1376 }
1377 sums
1378 }
1379 }
1380 } else {
1381 Vec::new()
1382 };
1383
1384 detect
1385 .par_iter()
1386 .enumerate()
1387 .map(|(i, det)| {
1388 let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
1389 let (x0, y0, x1, y1, roi_w, roi_h) = bbox_to_proto_roi(det, proto_w, proto_h);
1390
1391 let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
1393 let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
1394
1395 let mut mask_buf = vec![0u8; roi_h * roi_w];
1396
1397 match layout {
1398 edgefirst_tensor::ProtoLayout::Nhwc => {
1399 let stride_y = proto_w * num_protos;
1400 #[cfg(target_arch = "aarch64")]
1401 {
1402 for ly in 0..roi_h {
1403 let py = y0 + ly;
1404 let row_base = py * stride_y + x0 * num_protos;
1405 for lx in 0..roi_w {
1406 let pix_base = row_base + lx * num_protos;
1407 let proto_px = &protos[pix_base..pix_base + num_protos];
1408 let raw_dot = unsafe {
1409 dot_i16_i8_neon(coeff.as_ptr(), proto_px.as_ptr(), num_protos)
1410 };
1411 let correction = if zp_c != 0 {
1412 zp_c * proto_sums[py * proto_w + x0 + lx]
1413 } else {
1414 0
1415 };
1416 let logit = raw_dot - correction - bias;
1417 if logit > 0 {
1418 mask_buf[ly * roi_w + lx] = 255;
1419 }
1420 }
1421 }
1422 }
1423 #[cfg(not(target_arch = "aarch64"))]
1424 {
1425 for ly in 0..roi_h {
1426 let py = y0 + ly;
1427 let row_base = py * stride_y + x0 * num_protos;
1428 for lx in 0..roi_w {
1429 let pix_base = row_base + lx * num_protos;
1430 let proto_px = &protos[pix_base..pix_base + num_protos];
1431 let raw_dot = dot_i16_i8_scalar(coeff, proto_px, num_protos);
1432 let correction = if zp_c != 0 {
1433 zp_c * proto_sums[py * proto_w + x0 + lx]
1434 } else {
1435 0
1436 };
1437 let logit = raw_dot - correction - bias;
1438 if logit > 0 {
1439 mask_buf[ly * roi_w + lx] = 255;
1440 }
1441 }
1442 }
1443 }
1444 }
1445 edgefirst_tensor::ProtoLayout::Nchw => {
1446 let mut accum = vec![0i32; roi_h * roi_w];
1450 for c in 0..num_protos {
1451 let plane = &protos[c * hw..];
1452 let coeff_c = coeff[c] as i32;
1453 for ly in 0..roi_h {
1454 let py = y0 + ly;
1455 let row_start = py * proto_w + x0;
1456 let out_row_start = ly * roi_w;
1457 for lx in 0..roi_w {
1458 accum[out_row_start + lx] += coeff_c * plane[row_start + lx] as i32;
1459 }
1460 }
1461 }
1462 for ly in 0..roi_h {
1464 let py = y0 + ly;
1465 for lx in 0..roi_w {
1466 let idx = ly * roi_w + lx;
1467 let correction = if zp_c != 0 {
1468 zp_c * proto_sums[py * proto_w + x0 + lx]
1469 } else {
1470 0
1471 };
1472 let logit = accum[idx] - correction - bias;
1473 if logit > 0 {
1474 mask_buf[idx] = 255;
1475 }
1476 }
1477 }
1478 }
1479 }
1480
1481 let mask = ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1482 .expect("mask_buf length matches roi_h * roi_w");
1483 seg_from_roi(
1484 mask, x0, y0, x1, y1, proto_w, proto_h, lx0, inv_lw, ly0, inv_lh,
1485 )
1486 })
1487 .collect()
1488}
1489
1490#[allow(clippy::too_many_arguments)]
1501fn fused_dot_sign_f32_slice(
1502 protos: &[f32],
1503 coeff: &[f32],
1504 _proto_h: usize,
1505 proto_w: usize,
1506 y0: usize,
1507 x0: usize,
1508 roi_h: usize,
1509 roi_w: usize,
1510 num_protos: usize,
1511) -> ndarray::Array3<u8> {
1512 let stride_y = proto_w * num_protos;
1513 let mut mask_buf = vec![0u8; roi_h * roi_w];
1514 for y in 0..roi_h {
1515 let row_base = (y0 + y) * stride_y + x0 * num_protos;
1516 let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1517 for (x, out_px) in out_row.iter_mut().enumerate() {
1518 let base = row_base + x * num_protos;
1519 let mut acc = 0.0_f32;
1520 let mut k = 0;
1521 let chunks = num_protos / 4;
1522 for _ in 0..chunks {
1523 acc += coeff[k] * protos[base + k]
1524 + coeff[k + 1] * protos[base + k + 1]
1525 + coeff[k + 2] * protos[base + k + 2]
1526 + coeff[k + 3] * protos[base + k + 3];
1527 k += 4;
1528 }
1529 while k < num_protos {
1530 acc += coeff[k] * protos[base + k];
1531 k += 1;
1532 }
1533 if acc > 0.0 {
1534 *out_px = 255;
1535 }
1536 }
1537 }
1538 ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1539 .expect("mask_buf length matches roi_h * roi_w")
1540}
1541
1542#[allow(clippy::too_many_arguments)]
1551fn fused_dot_sign_f16_slice(
1552 protos: &[half::f16],
1553 coeff: &[f32],
1554 _proto_h: usize,
1555 proto_w: usize,
1556 y0: usize,
1557 x0: usize,
1558 roi_h: usize,
1559 roi_w: usize,
1560 num_protos: usize,
1561) -> ndarray::Array3<u8> {
1562 #[cfg(all(
1563 target_arch = "x86_64",
1564 target_feature = "f16c",
1565 target_feature = "fma"
1566 ))]
1567 {
1568 unsafe {
1570 fused_dot_sign_f16_slice_f16c(protos, coeff, proto_w, y0, x0, roi_h, roi_w, num_protos)
1571 }
1572 }
1573 #[cfg(not(all(
1574 target_arch = "x86_64",
1575 target_feature = "f16c",
1576 target_feature = "fma"
1577 )))]
1578 {
1579 fused_dot_sign_f16_slice_scalar(protos, coeff, proto_w, y0, x0, roi_h, roi_w, num_protos)
1580 }
1581}
1582
1583#[allow(clippy::too_many_arguments)]
1585fn fused_dot_sign_f16_slice_scalar(
1586 protos: &[half::f16],
1587 coeff: &[f32],
1588 proto_w: usize,
1589 y0: usize,
1590 x0: usize,
1591 roi_h: usize,
1592 roi_w: usize,
1593 num_protos: usize,
1594) -> ndarray::Array3<u8> {
1595 let stride_y = proto_w * num_protos;
1596 let mut mask_buf = vec![0u8; roi_h * roi_w];
1597 for y in 0..roi_h {
1598 let row_base = (y0 + y) * stride_y + x0 * num_protos;
1599 let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1600 for (x, out_px) in out_row.iter_mut().enumerate() {
1601 let base = row_base + x * num_protos;
1602 let mut acc = 0.0_f32;
1603 let mut k = 0;
1604 let chunks = num_protos / 4;
1605 for _ in 0..chunks {
1606 acc += coeff[k] * protos[base + k].to_f32()
1607 + coeff[k + 1] * protos[base + k + 1].to_f32()
1608 + coeff[k + 2] * protos[base + k + 2].to_f32()
1609 + coeff[k + 3] * protos[base + k + 3].to_f32();
1610 k += 4;
1611 }
1612 while k < num_protos {
1613 acc += coeff[k] * protos[base + k].to_f32();
1614 k += 1;
1615 }
1616 if acc > 0.0 {
1617 *out_px = 255;
1618 }
1619 }
1620 }
1621 ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1622 .expect("mask_buf length matches roi_h * roi_w")
1623}
1624
1625#[cfg(all(
1635 target_arch = "x86_64",
1636 target_feature = "f16c",
1637 target_feature = "fma"
1638))]
1639#[allow(clippy::too_many_arguments)]
1640#[target_feature(enable = "f16c,fma,avx")]
1641unsafe fn fused_dot_sign_f16_slice_f16c(
1642 protos: &[half::f16],
1643 coeff: &[f32],
1644 proto_w: usize,
1645 y0: usize,
1646 x0: usize,
1647 roi_h: usize,
1648 roi_w: usize,
1649 num_protos: usize,
1650) -> ndarray::Array3<u8> {
1651 unsafe {
1652 use core::arch::x86_64::{
1653 _mm256_castps256_ps128, _mm256_cvtph_ps, _mm256_extractf128_ps, _mm256_fmadd_ps,
1654 _mm256_loadu_ps, _mm256_setzero_ps, _mm_add_ps, _mm_cvtss_f32, _mm_hadd_ps,
1655 _mm_loadu_si128,
1656 };
1657
1658 let stride_y = proto_w * num_protos;
1659 let chunks8 = num_protos / 8;
1660 let mut mask_buf = vec![0u8; roi_h * roi_w];
1661
1662 for y in 0..roi_h {
1663 let row_base = (y0 + y) * stride_y + x0 * num_protos;
1664 let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1665 for (x, out_px) in out_row.iter_mut().enumerate() {
1666 let base = row_base + x * num_protos;
1667 let mut acc_v = _mm256_setzero_ps();
1668 let mut k = 0;
1669 for _ in 0..chunks8 {
1670 let p_ptr = protos
1671 .as_ptr()
1672 .add(base + k)
1673 .cast::<core::arch::x86_64::__m128i>();
1674 let raw = _mm_loadu_si128(p_ptr);
1675 let widened = _mm256_cvtph_ps(raw);
1676 let coeffs_v = _mm256_loadu_ps(coeff.as_ptr().add(k));
1677 acc_v = _mm256_fmadd_ps(coeffs_v, widened, acc_v);
1678 k += 8;
1679 }
1680 let lo = _mm256_castps256_ps128(acc_v);
1682 let hi = _mm256_extractf128_ps::<1>(acc_v);
1683 let sum4 = _mm_add_ps(lo, hi);
1684 let sum2 = _mm_hadd_ps(sum4, sum4);
1685 let sum1 = _mm_hadd_ps(sum2, sum2);
1686 let mut acc = _mm_cvtss_f32(sum1);
1687
1688 while k < num_protos {
1690 acc += coeff[k] * protos[base + k].to_f32();
1691 k += 1;
1692 }
1693
1694 if acc > 0.0 {
1695 *out_px = 255;
1696 }
1697 }
1698 }
1699 ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1700 .expect("mask_buf length matches roi_h * roi_w")
1701 }
1702}
1703
1704#[allow(clippy::too_many_arguments)]
1708fn fused_dequant_dot_sign_i8_slice(
1709 protos: &[i8],
1710 coeff: &[f32],
1711 quant: &edgefirst_tensor::Quantization,
1712 _proto_h: usize,
1713 proto_w: usize,
1714 y0: usize,
1715 x0: usize,
1716 roi_h: usize,
1717 roi_w: usize,
1718 num_protos: usize,
1719) -> crate::Result<ndarray::Array3<u8>> {
1720 use edgefirst_tensor::QuantMode;
1721 let stride_y = proto_w * num_protos;
1722
1723 let mut stack_scratch = [0.0_f32; 64];
1725 let mut heap_scratch: Vec<f32>;
1726 let scaled_coeff: &mut [f32] = if num_protos <= stack_scratch.len() {
1727 &mut stack_scratch[..num_protos]
1728 } else {
1729 heap_scratch = vec![0.0_f32; num_protos];
1730 heap_scratch.as_mut_slice()
1731 };
1732 let zp_offset: f32 = match quant.mode() {
1733 QuantMode::PerTensorSymmetric { scale } => {
1734 for k in 0..num_protos {
1735 scaled_coeff[k] = coeff[k] * scale;
1736 }
1737 0.0
1738 }
1739 QuantMode::PerTensor { scale, zero_point } => {
1740 for k in 0..num_protos {
1741 scaled_coeff[k] = coeff[k] * scale;
1742 }
1743 zero_point as f32 * scaled_coeff.iter().take(num_protos).sum::<f32>()
1744 }
1745 QuantMode::PerChannelSymmetric { scales, axis } => {
1746 if axis != 2 {
1747 return Err(crate::Error::NotSupported(format!(
1748 "per-channel quantization on axis {axis} not supported \
1749 (only channel axis 2 is implemented on this kernel)"
1750 )));
1751 }
1752 for k in 0..num_protos {
1753 scaled_coeff[k] = coeff[k] * scales[k];
1754 }
1755 0.0
1756 }
1757 QuantMode::PerChannel {
1758 scales,
1759 zero_points,
1760 axis,
1761 } => {
1762 if axis != 2 {
1763 return Err(crate::Error::NotSupported(format!(
1764 "per-channel quantization on axis {axis} not supported \
1765 (only channel axis 2 is implemented on this kernel)"
1766 )));
1767 }
1768 for k in 0..num_protos {
1769 scaled_coeff[k] = coeff[k] * scales[k];
1770 }
1771 (0..num_protos)
1772 .map(|k| scaled_coeff[k] * zero_points[k] as f32)
1773 .sum()
1774 }
1775 };
1776
1777 let mut mask_buf = vec![0u8; roi_h * roi_w];
1778 for y in 0..roi_h {
1779 let row_base = (y0 + y) * stride_y + (x0) * num_protos;
1780 let out_row = &mut mask_buf[y * roi_w..(y + 1) * roi_w];
1781 for (x, out_px) in out_row.iter_mut().enumerate() {
1782 let base = row_base + x * num_protos;
1783 let mut acc = 0.0_f32;
1784 let mut k = 0;
1785 let chunks = num_protos / 4;
1786 for _ in 0..chunks {
1787 let p0 = protos[base + k] as f32;
1788 let p1 = protos[base + k + 1] as f32;
1789 let p2 = protos[base + k + 2] as f32;
1790 let p3 = protos[base + k + 3] as f32;
1791 acc += scaled_coeff[k] * p0
1792 + scaled_coeff[k + 1] * p1
1793 + scaled_coeff[k + 2] * p2
1794 + scaled_coeff[k + 3] * p3;
1795 k += 4;
1796 }
1797 while k < num_protos {
1798 acc += scaled_coeff[k] * protos[base + k] as f32;
1799 k += 1;
1800 }
1801 if acc > zp_offset {
1802 *out_px = 255;
1803 }
1804 }
1805 }
1806 Ok(ndarray::Array3::from_shape_vec((roi_h, roi_w, 1), mask_buf)
1807 .expect("mask_buf length matches roi_h * roi_w"))
1808}
1809
1810#[allow(clippy::too_many_arguments)]
1811fn scaled_segmentations_f32_slice(
1812 detect: &[crate::DetectBox],
1813 coeff_all: &[f32],
1814 protos: &[f32],
1815 proto_h: usize,
1816 proto_w: usize,
1817 num_protos: usize,
1818 letterbox: Option<[f32; 4]>,
1819 width: u32,
1820 height: u32,
1821) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
1822 scaled_run(
1823 detect,
1824 coeff_all,
1825 protos,
1826 proto_h,
1827 proto_w,
1828 num_protos,
1829 letterbox,
1830 width,
1831 height,
1832 1.0,
1833 |p, _| *p,
1834 )
1835}
1836
1837#[allow(clippy::too_many_arguments)]
1838fn scaled_segmentations_f16_slice(
1839 detect: &[crate::DetectBox],
1840 coeff_all: &[f32],
1841 protos: &[half::f16],
1842 proto_h: usize,
1843 proto_w: usize,
1844 num_protos: usize,
1845 letterbox: Option<[f32; 4]>,
1846 width: u32,
1847 height: u32,
1848) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
1849 scaled_run(
1850 detect,
1851 coeff_all,
1852 protos,
1853 proto_h,
1854 proto_w,
1855 num_protos,
1856 letterbox,
1857 width,
1858 height,
1859 1.0,
1860 |p: &half::f16, _| p.to_f32(),
1861 )
1862}
1863
1864#[allow(clippy::too_many_arguments)]
1865fn scaled_segmentations_i8_slice(
1866 detect: &[crate::DetectBox],
1867 coeff_all: &[f32],
1868 protos: &[i8],
1869 proto_h: usize,
1870 proto_w: usize,
1871 num_protos: usize,
1872 quant: &edgefirst_tensor::Quantization,
1873 letterbox: Option<[f32; 4]>,
1874 width: u32,
1875 height: u32,
1876) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
1877 use edgefirst_tensor::QuantMode;
1878 let (scale, zp) = match quant.mode() {
1882 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1883 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1884 QuantMode::PerChannel { axis, .. } | QuantMode::PerChannelSymmetric { axis, .. } => {
1885 return Err(crate::Error::NotSupported(format!(
1886 "per-channel quantization (axis={axis}) on scaled seg path \
1887 not yet supported"
1888 )));
1889 }
1890 };
1891 scaled_run(
1892 detect,
1893 coeff_all,
1894 protos,
1895 proto_h,
1896 proto_w,
1897 num_protos,
1898 letterbox,
1899 width,
1900 height,
1901 scale,
1902 move |p: &i8, _| *p as f32 - zp,
1903 )
1904}
1905
1906#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
1922#[inline(always)]
1923fn dot_i8_scalar(coeff: &[i8], proto: &[i8], n: usize) -> i32 {
1924 let mut acc: i32 = 0;
1925 let chunks = n / 4;
1926 let mut k = 0;
1927 for _ in 0..chunks {
1928 acc += coeff[k] as i32 * proto[k] as i32
1929 + coeff[k + 1] as i32 * proto[k + 1] as i32
1930 + coeff[k + 2] as i32 * proto[k + 2] as i32
1931 + coeff[k + 3] as i32 * proto[k + 3] as i32;
1932 k += 4;
1933 }
1934 while k < n {
1935 acc += coeff[k] as i32 * proto[k] as i32;
1936 k += 1;
1937 }
1938 acc
1939}
1940
1941#[cfg(target_arch = "aarch64")]
1943#[inline(always)]
1944unsafe fn dot_i8_neon_base(coeff: *const i8, proto: *const i8, n: usize) -> i32 {
1945 unsafe {
1946 use std::arch::aarch64::*; let mut acc = vdupq_n_s32(0);
1948 let full_chunks = n / 16;
1949 let mut offset = 0usize;
1950 for _ in 0..full_chunks {
1951 let c = vld1q_s8(coeff.add(offset));
1952 let p = vld1q_s8(proto.add(offset));
1953 let lo = vmull_s8(vget_low_s8(c), vget_low_s8(p));
1955 let hi = vmull_high_s8(c, p);
1956 acc = vpadalq_s16(acc, lo);
1957 acc = vpadalq_s16(acc, hi);
1958 offset += 16;
1959 }
1960 let remainder = n - offset;
1962 if remainder >= 8 {
1963 let c = vld1_s8(coeff.add(offset));
1964 let p = vld1_s8(proto.add(offset));
1965 let prod = vmull_s8(c, p);
1966 acc = vpadalq_s16(acc, prod);
1967 offset += 8;
1968 }
1969 let mut scalar_acc = vaddvq_s32(acc);
1970 while offset < n {
1971 scalar_acc += *coeff.add(offset) as i32 * *proto.add(offset) as i32;
1972 offset += 1;
1973 }
1974 scalar_acc
1975 }
1976}
1977
1978#[cfg(target_arch = "aarch64")]
1982#[inline(always)]
1983unsafe fn dot_i8_neon_dotprod(coeff: *const i8, proto: *const i8, n: usize) -> i32 {
1984 unsafe {
1985 use std::arch::aarch64::*; let mut acc = vdupq_n_s32(0);
1987 let full_chunks = n / 16;
1988 let mut offset = 0usize;
1989 for _ in 0..full_chunks {
1990 let c = vld1q_s8(coeff.add(offset));
1991 let p = vld1q_s8(proto.add(offset));
1992 let result: int32x4_t;
1996 core::arch::asm!(
1997 ".arch_extension dotprod",
1998 "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
1999 acc = inout(vreg) acc => result,
2000 a = in(vreg) c,
2001 b = in(vreg) p,
2002 options(pure, nomem, nostack),
2003 );
2004 acc = result;
2005 offset += 16;
2006 }
2007 let mut scalar_acc = vaddvq_s32(acc);
2008 while offset < n {
2010 scalar_acc += *coeff.add(offset) as i32 * *proto.add(offset) as i32;
2011 offset += 1;
2012 }
2013 scalar_acc
2014 }
2015}
2016
2017#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
2019#[inline(always)]
2020fn dot_i16_i8_scalar(coeff: &[i16], proto: &[i8], n: usize) -> i32 {
2021 let mut acc: i32 = 0;
2022 let chunks = n / 4;
2023 let mut k = 0;
2024 for _ in 0..chunks {
2025 acc += coeff[k] as i32 * proto[k] as i32
2026 + coeff[k + 1] as i32 * proto[k + 1] as i32
2027 + coeff[k + 2] as i32 * proto[k + 2] as i32
2028 + coeff[k + 3] as i32 * proto[k + 3] as i32;
2029 k += 4;
2030 }
2031 while k < n {
2032 acc += coeff[k] as i32 * proto[k] as i32;
2033 k += 1;
2034 }
2035 acc
2036}
2037
2038#[cfg(target_arch = "aarch64")]
2041#[inline(always)]
2042unsafe fn dot_i16_i8_neon(coeff: *const i16, proto: *const i8, n: usize) -> i32 {
2043 unsafe {
2044 use std::arch::aarch64::*; let mut acc = vdupq_n_s32(0);
2046 let full_chunks = n / 8;
2047 let mut offset = 0usize;
2048 for _ in 0..full_chunks {
2049 let c = vld1q_s16(coeff.add(offset));
2050 let p_raw = vld1_s8(proto.add(offset));
2051 let p = vmovl_s8(p_raw);
2052 acc = vmlal_s16(acc, vget_low_s16(c), vget_low_s16(p));
2053 acc = vmlal_high_s16(acc, c, p);
2054 offset += 8;
2055 }
2056 let mut scalar_acc = vaddvq_s32(acc);
2057 while offset < n {
2058 scalar_acc += *coeff.add(offset) as i32 * *proto.add(offset) as i32;
2059 offset += 1;
2060 }
2061 scalar_acc
2062 }
2063}
2064
2065#[cfg(target_arch = "aarch64")]
2068#[inline(always)]
2069#[allow(clippy::too_many_arguments)]
2070fn compute_logits_dotprod(
2071 logits: &mut [i32],
2072 coeff: &[i8],
2073 protos: &[i8],
2074 proto_sums: &[i32],
2075 proto_w: usize,
2076 proto_x0: usize,
2077 proto_y0: usize,
2078 roi_w: usize,
2079 roi_h: usize,
2080 stride_y: usize,
2081 num_protos: usize,
2082 zp_c: i32,
2083 bias: i32,
2084) {
2085 for ly_idx in 0..roi_h {
2086 let py = proto_y0 + ly_idx;
2087 let row_base = py * stride_y + proto_x0 * num_protos;
2088 for lx_idx in 0..roi_w {
2089 let pix_base = row_base + lx_idx * num_protos;
2090 let proto_px = &protos[pix_base..pix_base + num_protos];
2091 let raw_dot =
2092 unsafe { dot_i8_neon_dotprod(coeff.as_ptr(), proto_px.as_ptr(), num_protos) };
2093 let correction = if zp_c != 0 {
2094 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2095 } else {
2096 0
2097 };
2098 logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2099 }
2100 }
2101}
2102
2103#[cfg(target_arch = "aarch64")]
2106#[inline(always)]
2107#[allow(clippy::too_many_arguments)]
2108fn compute_logits_base(
2109 logits: &mut [i32],
2110 coeff: &[i8],
2111 protos: &[i8],
2112 proto_sums: &[i32],
2113 proto_w: usize,
2114 proto_x0: usize,
2115 proto_y0: usize,
2116 roi_w: usize,
2117 roi_h: usize,
2118 stride_y: usize,
2119 num_protos: usize,
2120 zp_c: i32,
2121 bias: i32,
2122) {
2123 for ly_idx in 0..roi_h {
2124 let py = proto_y0 + ly_idx;
2125 let row_base = py * stride_y + proto_x0 * num_protos;
2126 for lx_idx in 0..roi_w {
2127 let pix_base = row_base + lx_idx * num_protos;
2128 let proto_px = &protos[pix_base..pix_base + num_protos];
2129 let raw_dot =
2130 unsafe { dot_i8_neon_base(coeff.as_ptr(), proto_px.as_ptr(), num_protos) };
2131 let correction = if zp_c != 0 {
2132 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2133 } else {
2134 0
2135 };
2136 logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2137 }
2138 }
2139}
2140
2141#[allow(clippy::too_many_arguments)]
2142fn scaled_segmentations_i8_i8(
2143 detect: &[crate::DetectBox],
2144 coeff_all: &[i8],
2145 coeff_quant: &edgefirst_tensor::Quantization,
2146 protos: &[i8],
2147 proto_quant: &edgefirst_tensor::Quantization,
2148 proto_h: usize,
2149 proto_w: usize,
2150 num_protos: usize,
2151 letterbox: Option<[f32; 4]>,
2152 width: u32,
2153 height: u32,
2154 layout: edgefirst_tensor::ProtoLayout,
2155) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
2156 use edgefirst_tensor::QuantMode;
2157
2158 let _span = tracing::trace_span!(
2159 "image.materialize_masks.kernel_i8_scaled",
2160 n = detect.len(),
2161 proto_h,
2162 proto_w,
2163 num_protos,
2164 width,
2165 height,
2166 ?layout,
2167 )
2168 .entered();
2169
2170 let zp_c: i32 = match coeff_quant.mode() {
2171 QuantMode::PerTensor { zero_point, .. } => zero_point,
2172 QuantMode::PerTensorSymmetric { .. } => 0,
2173 _ => {
2174 return Err(crate::Error::NotSupported(
2175 "per-channel coeff quantization not supported".into(),
2176 ))
2177 }
2178 };
2179 let zp_p: i32 = match proto_quant.mode() {
2180 QuantMode::PerTensor { zero_point, .. } => zero_point,
2181 QuantMode::PerTensorSymmetric { .. } => 0,
2182 _ => {
2183 return Err(crate::Error::NotSupported(
2184 "per-channel proto quantization not supported".into(),
2185 ))
2186 }
2187 };
2188
2189 let (lx0, lw, ly0, lh) = match letterbox {
2190 Some([lx0, ly0, lx1, ly1]) => {
2191 let lw = if lx1 > lx0 { lx1 - lx0 } else { 1.0 };
2199 let lh = if ly1 > ly0 { ly1 - ly0 } else { 1.0 };
2200 (lx0, lw, ly0, lh)
2201 }
2202 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
2203 };
2204 let out_w = width as usize;
2205 let out_h = height as usize;
2206 let hw = proto_h * proto_w;
2207
2208 let proto_sums: Vec<i32> = if zp_c != 0 {
2210 match layout {
2211 edgefirst_tensor::ProtoLayout::Nhwc => (0..hw)
2212 .map(|px_idx| {
2213 let base = px_idx * num_protos;
2214 let mut s: i32 = 0;
2215 for k in 0..num_protos {
2216 s += protos[base + k] as i32;
2217 }
2218 s
2219 })
2220 .collect(),
2221 edgefirst_tensor::ProtoLayout::Nchw => {
2222 let mut sums = vec![0i32; hw];
2223 for c in 0..num_protos {
2224 let plane = &protos[c * hw..];
2225 for (px, s) in sums.iter_mut().enumerate() {
2226 *s += plane[px] as i32;
2227 }
2228 }
2229 sums
2230 }
2231 }
2232 } else {
2233 Vec::new()
2234 };
2235
2236 #[cfg(target_arch = "aarch64")]
2238 let use_dotprod = std::arch::is_aarch64_feature_detected!("dotprod");
2239
2240 let stride_y = proto_w * num_protos;
2242
2243 detect
2244 .par_iter()
2245 .enumerate()
2246 .map(|(i, det)| {
2247 let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
2248 let bbox = det.bbox.to_canonical();
2249 let xmin = ((bbox.xmin - lx0) / lw).clamp(0.0, 1.0);
2250 let ymin = ((bbox.ymin - ly0) / lh).clamp(0.0, 1.0);
2251 let xmax = ((bbox.xmax - lx0) / lw).clamp(0.0, 1.0);
2252 let ymax = ((bbox.ymax - ly0) / lh).clamp(0.0, 1.0);
2253 let px0 = (xmin * out_w as f32).round() as usize;
2254 let py0 = (ymin * out_h as f32).round() as usize;
2255 let px1 = ((xmax * out_w as f32).round() as usize).min(out_w);
2256 let py1 = ((ymax * out_h as f32).round() as usize).min(out_h);
2257 let bbox_w = px1.saturating_sub(px0).max(1);
2258 let bbox_h = py1.saturating_sub(py0).max(1);
2259
2260 let sample_x_at = |px: f32| -> f32 {
2262 let model_x_norm = lx0 + (px + 0.5) / out_w as f32 * lw;
2263 model_x_norm * proto_w as f32 - 0.5
2264 };
2265 let sample_y_at = |py: f32| -> f32 {
2266 let model_y_norm = ly0 + (py + 0.5) / out_h as f32 * lh;
2267 model_y_norm * proto_h as f32 - 0.5
2268 };
2269 let s_x_min = sample_x_at(px0 as f32);
2270 let s_x_max = sample_x_at((px1 as f32) - 1.0);
2271 let s_y_min = sample_y_at(py0 as f32);
2272 let s_y_max = sample_y_at((py1 as f32) - 1.0);
2273 let proto_x0 = (s_x_min.floor() as isize)
2274 .max(0)
2275 .min(proto_w.saturating_sub(1) as isize) as usize;
2276 let proto_x1 = ((s_x_max.ceil() as isize) + 1).max(0).min(proto_w as isize) as usize;
2277 let proto_y0 = (s_y_min.floor() as isize)
2278 .max(0)
2279 .min(proto_h.saturating_sub(1) as isize) as usize;
2280 let proto_y1 = ((s_y_max.ceil() as isize) + 1).max(0).min(proto_h as isize) as usize;
2281 let roi_w = proto_x1.saturating_sub(proto_x0).max(1);
2282 let roi_h = proto_y1.saturating_sub(proto_y0).max(1);
2283
2284 let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
2286 let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
2287
2288 let mut logits = vec![0_i32; roi_h * roi_w];
2290 match layout {
2291 edgefirst_tensor::ProtoLayout::Nhwc => {
2292 #[cfg(target_arch = "aarch64")]
2293 {
2294 if use_dotprod {
2295 compute_logits_dotprod(
2296 &mut logits,
2297 coeff,
2298 protos,
2299 &proto_sums,
2300 proto_w,
2301 proto_x0,
2302 proto_y0,
2303 roi_w,
2304 roi_h,
2305 stride_y,
2306 num_protos,
2307 zp_c,
2308 bias,
2309 );
2310 } else {
2311 compute_logits_base(
2312 &mut logits,
2313 coeff,
2314 protos,
2315 &proto_sums,
2316 proto_w,
2317 proto_x0,
2318 proto_y0,
2319 roi_w,
2320 roi_h,
2321 stride_y,
2322 num_protos,
2323 zp_c,
2324 bias,
2325 );
2326 }
2327 }
2328 #[cfg(not(target_arch = "aarch64"))]
2329 {
2330 for ly_idx in 0..roi_h {
2331 let py = proto_y0 + ly_idx;
2332 let row_base = py * stride_y + proto_x0 * num_protos;
2333 for lx_idx in 0..roi_w {
2334 let pix_base = row_base + lx_idx * num_protos;
2335 let proto_px = &protos[pix_base..pix_base + num_protos];
2336 let raw_dot = dot_i8_scalar(coeff, proto_px, num_protos);
2337 let correction = if zp_c != 0 {
2338 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2339 } else {
2340 0
2341 };
2342 logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2343 }
2344 }
2345 }
2346 }
2347 edgefirst_tensor::ProtoLayout::Nchw => {
2348 for c in 0..num_protos {
2350 let plane = &protos[c * hw..];
2351 let coeff_c = coeff[c] as i32;
2352 for ly_idx in 0..roi_h {
2353 let py = proto_y0 + ly_idx;
2354 let row_start = py * proto_w + proto_x0;
2355 let out_row_start = ly_idx * roi_w;
2356 for lx_idx in 0..roi_w {
2357 logits[out_row_start + lx_idx] +=
2358 coeff_c * plane[row_start + lx_idx] as i32;
2359 }
2360 }
2361 }
2362 for ly_idx in 0..roi_h {
2364 let py = proto_y0 + ly_idx;
2365 for lx_idx in 0..roi_w {
2366 let idx = ly_idx * roi_w + lx_idx;
2367 let correction = if zp_c != 0 {
2368 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2369 } else {
2370 0
2371 };
2372 logits[idx] -= correction + bias;
2373 }
2374 }
2375 }
2376 }
2377
2378 let roi_last_x = roi_w.saturating_sub(1);
2381 let roi_last_y = roi_h.saturating_sub(1);
2382
2383 const FRAC_BITS: i32 = 10;
2385 const FRAC_SCALE: i32 = 1 << FRAC_BITS; let x_coords: Vec<(usize, usize, i32)> = (0..bbox_w)
2387 .map(|xi| {
2388 let sample_x = sample_x_at((px0 + xi) as f32) - proto_x0 as f32;
2389 let x_floor = sample_x.floor();
2390 let x_lo = (x_floor as isize).max(0).min(roi_last_x as isize) as usize;
2391 let x_hi = (x_lo + 1).min(roi_w - 1);
2392 let x_frac = ((sample_x - x_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2393 (x_lo, x_hi, x_frac)
2394 })
2395 .collect();
2396
2397 let mut tile_buf = vec![0u8; bbox_h * bbox_w];
2398 for yi in 0..bbox_h {
2399 let sample_y = sample_y_at((py0 + yi) as f32) - proto_y0 as f32;
2400 let y_floor = sample_y.floor();
2401 let y_lo = (y_floor as isize).max(0).min(roi_last_y as isize) as usize;
2402 let y_hi = (y_lo + 1).min(roi_h - 1);
2403 let y_frac = ((sample_y - y_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2404 let y_frac_inv = FRAC_SCALE - y_frac;
2405 let row_lo = &logits[y_lo * roi_w..y_lo * roi_w + roi_w];
2406 let row_hi = &logits[y_hi * roi_w..y_hi * roi_w + roi_w];
2407 let out_row = &mut tile_buf[yi * bbox_w..(yi + 1) * bbox_w];
2408
2409 for (xi, &(x_lo, x_hi, x_frac)) in x_coords.iter().enumerate() {
2410 let tl = row_lo[x_lo];
2411 let tr = row_lo[x_hi];
2412 let bl = row_hi[x_lo];
2413 let br = row_hi[x_hi];
2414
2415 if (tl & tr & bl & br) < 0 {
2419 continue;
2421 }
2422 if tl > 0 && tr > 0 && bl > 0 && br > 0 {
2423 out_row[xi] = 255;
2425 continue;
2426 }
2427
2428 let x_frac_inv = FRAC_SCALE - x_frac;
2430 let l0 = tl as i64 * x_frac_inv as i64 + tr as i64 * x_frac as i64;
2431 let l1 = bl as i64 * x_frac_inv as i64 + br as i64 * x_frac as i64;
2432 let logit = l0 * y_frac_inv as i64 + l1 * y_frac as i64;
2433 out_row[xi] = if logit > 0 { 255 } else { 0 };
2434 }
2435 }
2436
2437 let tile = ndarray::Array3::from_shape_vec((bbox_h, bbox_w, 1), tile_buf)
2438 .expect("tile_buf length matches bbox_h * bbox_w");
2439 Ok(edgefirst_tensor::Segmentation {
2440 xmin,
2441 ymin,
2442 xmax,
2443 ymax,
2444 segmentation: crate::cpu::masks::mask_to_tensor(tile.view())?,
2445 })
2446 })
2447 .collect()
2448}
2449
2450#[allow(clippy::too_many_arguments)]
2451fn scaled_segmentations_i16_i8(
2452 detect: &[crate::DetectBox],
2453 coeff_all: &[i16],
2454 coeff_quant: &edgefirst_tensor::Quantization,
2455 protos: &[i8],
2456 proto_quant: &edgefirst_tensor::Quantization,
2457 proto_h: usize,
2458 proto_w: usize,
2459 num_protos: usize,
2460 letterbox: Option<[f32; 4]>,
2461 width: u32,
2462 height: u32,
2463 layout: edgefirst_tensor::ProtoLayout,
2464) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
2465 use edgefirst_tensor::QuantMode;
2466
2467 let _span = tracing::trace_span!(
2468 "image.materialize_masks.kernel_i16xi8_scaled",
2469 n = detect.len(),
2470 proto_h,
2471 proto_w,
2472 num_protos,
2473 width,
2474 height,
2475 ?layout,
2476 )
2477 .entered();
2478
2479 let zp_c: i32 = match coeff_quant.mode() {
2480 QuantMode::PerTensor { zero_point, .. } => zero_point,
2481 QuantMode::PerTensorSymmetric { .. } => 0,
2482 _ => {
2483 return Err(crate::Error::NotSupported(
2484 "per-channel coeff quantization not supported".into(),
2485 ))
2486 }
2487 };
2488 let zp_p: i32 = match proto_quant.mode() {
2489 QuantMode::PerTensor { zero_point, .. } => zero_point,
2490 QuantMode::PerTensorSymmetric { .. } => 0,
2491 _ => {
2492 return Err(crate::Error::NotSupported(
2493 "per-channel proto quantization not supported".into(),
2494 ))
2495 }
2496 };
2497
2498 let (lx0, lw, ly0, lh) = match letterbox {
2499 Some([lx0, ly0, lx1, ly1]) => {
2500 let lw = if lx1 > lx0 { lx1 - lx0 } else { 1.0 };
2508 let lh = if ly1 > ly0 { ly1 - ly0 } else { 1.0 };
2509 (lx0, lw, ly0, lh)
2510 }
2511 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
2512 };
2513 let out_w = width as usize;
2514 let out_h = height as usize;
2515 let hw = proto_h * proto_w;
2516
2517 let proto_sums: Vec<i32> = if zp_c != 0 {
2519 match layout {
2520 edgefirst_tensor::ProtoLayout::Nhwc => (0..hw)
2521 .map(|px_idx| {
2522 let base = px_idx * num_protos;
2523 let mut s: i32 = 0;
2524 for k in 0..num_protos {
2525 s += protos[base + k] as i32;
2526 }
2527 s
2528 })
2529 .collect(),
2530 edgefirst_tensor::ProtoLayout::Nchw => {
2531 let mut sums = vec![0i32; hw];
2532 for c in 0..num_protos {
2533 let plane = &protos[c * hw..];
2534 for (px, s) in sums.iter_mut().enumerate() {
2535 *s += plane[px] as i32;
2536 }
2537 }
2538 sums
2539 }
2540 }
2541 } else {
2542 Vec::new()
2543 };
2544
2545 let stride_y = proto_w * num_protos;
2547
2548 detect
2549 .par_iter()
2550 .enumerate()
2551 .map(|(i, det)| {
2552 let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
2553 let bbox = det.bbox.to_canonical();
2554 let xmin = ((bbox.xmin - lx0) / lw).clamp(0.0, 1.0);
2555 let ymin = ((bbox.ymin - ly0) / lh).clamp(0.0, 1.0);
2556 let xmax = ((bbox.xmax - lx0) / lw).clamp(0.0, 1.0);
2557 let ymax = ((bbox.ymax - ly0) / lh).clamp(0.0, 1.0);
2558 let px0 = (xmin * out_w as f32).round() as usize;
2559 let py0 = (ymin * out_h as f32).round() as usize;
2560 let px1 = ((xmax * out_w as f32).round() as usize).min(out_w);
2561 let py1 = ((ymax * out_h as f32).round() as usize).min(out_h);
2562 let bbox_w = px1.saturating_sub(px0).max(1);
2563 let bbox_h = py1.saturating_sub(py0).max(1);
2564
2565 let sample_x_at = |px: f32| -> f32 {
2567 let model_x_norm = lx0 + (px + 0.5) / out_w as f32 * lw;
2568 model_x_norm * proto_w as f32 - 0.5
2569 };
2570 let sample_y_at = |py: f32| -> f32 {
2571 let model_y_norm = ly0 + (py + 0.5) / out_h as f32 * lh;
2572 model_y_norm * proto_h as f32 - 0.5
2573 };
2574 let s_x_min = sample_x_at(px0 as f32);
2575 let s_x_max = sample_x_at((px1 as f32) - 1.0);
2576 let s_y_min = sample_y_at(py0 as f32);
2577 let s_y_max = sample_y_at((py1 as f32) - 1.0);
2578 let proto_x0 = (s_x_min.floor() as isize)
2579 .max(0)
2580 .min(proto_w.saturating_sub(1) as isize) as usize;
2581 let proto_x1 = ((s_x_max.ceil() as isize) + 1).max(0).min(proto_w as isize) as usize;
2582 let proto_y0 = (s_y_min.floor() as isize)
2583 .max(0)
2584 .min(proto_h.saturating_sub(1) as isize) as usize;
2585 let proto_y1 = ((s_y_max.ceil() as isize) + 1).max(0).min(proto_h as isize) as usize;
2586 let roi_w = proto_x1.saturating_sub(proto_x0).max(1);
2587 let roi_h = proto_y1.saturating_sub(proto_y0).max(1);
2588
2589 let coeff_sum: i32 = coeff.iter().map(|&c| c as i32).sum();
2591 let bias = zp_p * coeff_sum - (num_protos as i32) * zp_c * zp_p;
2592
2593 let mut logits = vec![0_i32; roi_h * roi_w];
2595 match layout {
2596 edgefirst_tensor::ProtoLayout::Nhwc => {
2597 #[cfg(target_arch = "aarch64")]
2598 {
2599 for ly_idx in 0..roi_h {
2600 let py = proto_y0 + ly_idx;
2601 let row_base = py * stride_y + proto_x0 * num_protos;
2602 for lx_idx in 0..roi_w {
2603 let pix_base = row_base + lx_idx * num_protos;
2604 let proto_px = &protos[pix_base..pix_base + num_protos];
2605 let raw_dot = unsafe {
2606 dot_i16_i8_neon(coeff.as_ptr(), proto_px.as_ptr(), num_protos)
2607 };
2608 let correction = if zp_c != 0 {
2609 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2610 } else {
2611 0
2612 };
2613 logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2614 }
2615 }
2616 }
2617 #[cfg(not(target_arch = "aarch64"))]
2618 {
2619 for ly_idx in 0..roi_h {
2620 let py = proto_y0 + ly_idx;
2621 let row_base = py * stride_y + proto_x0 * num_protos;
2622 for lx_idx in 0..roi_w {
2623 let pix_base = row_base + lx_idx * num_protos;
2624 let proto_px = &protos[pix_base..pix_base + num_protos];
2625 let raw_dot = dot_i16_i8_scalar(coeff, proto_px, num_protos);
2626 let correction = if zp_c != 0 {
2627 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2628 } else {
2629 0
2630 };
2631 logits[ly_idx * roi_w + lx_idx] = raw_dot - correction - bias;
2632 }
2633 }
2634 }
2635 }
2636 edgefirst_tensor::ProtoLayout::Nchw => {
2637 for c in 0..num_protos {
2639 let plane = &protos[c * hw..];
2640 let coeff_c = coeff[c] as i32;
2641 for ly_idx in 0..roi_h {
2642 let py = proto_y0 + ly_idx;
2643 let row_start = py * proto_w + proto_x0;
2644 let out_row_start = ly_idx * roi_w;
2645 for lx_idx in 0..roi_w {
2646 logits[out_row_start + lx_idx] +=
2647 coeff_c * plane[row_start + lx_idx] as i32;
2648 }
2649 }
2650 }
2651 for ly_idx in 0..roi_h {
2653 let py = proto_y0 + ly_idx;
2654 for lx_idx in 0..roi_w {
2655 let idx = ly_idx * roi_w + lx_idx;
2656 let correction = if zp_c != 0 {
2657 zp_c * proto_sums[py * proto_w + proto_x0 + lx_idx]
2658 } else {
2659 0
2660 };
2661 logits[idx] -= correction + bias;
2662 }
2663 }
2664 }
2665 }
2666
2667 let roi_last_x = roi_w.saturating_sub(1);
2670 let roi_last_y = roi_h.saturating_sub(1);
2671
2672 const FRAC_BITS: i32 = 10;
2674 const FRAC_SCALE: i32 = 1 << FRAC_BITS; let x_coords: Vec<(usize, usize, i32)> = (0..bbox_w)
2676 .map(|xi| {
2677 let sample_x = sample_x_at((px0 + xi) as f32) - proto_x0 as f32;
2678 let x_floor = sample_x.floor();
2679 let x_lo = (x_floor as isize).max(0).min(roi_last_x as isize) as usize;
2680 let x_hi = (x_lo + 1).min(roi_w - 1);
2681 let x_frac = ((sample_x - x_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2682 (x_lo, x_hi, x_frac)
2683 })
2684 .collect();
2685
2686 let mut tile_buf = vec![0u8; bbox_h * bbox_w];
2687 for yi in 0..bbox_h {
2688 let sample_y = sample_y_at((py0 + yi) as f32) - proto_y0 as f32;
2689 let y_floor = sample_y.floor();
2690 let y_lo = (y_floor as isize).max(0).min(roi_last_y as isize) as usize;
2691 let y_hi = (y_lo + 1).min(roi_h - 1);
2692 let y_frac = ((sample_y - y_floor).clamp(0.0, 1.0) * FRAC_SCALE as f32) as i32;
2693 let y_frac_inv = FRAC_SCALE - y_frac;
2694 let row_lo = &logits[y_lo * roi_w..y_lo * roi_w + roi_w];
2695 let row_hi = &logits[y_hi * roi_w..y_hi * roi_w + roi_w];
2696 let out_row = &mut tile_buf[yi * bbox_w..(yi + 1) * bbox_w];
2697
2698 for (xi, &(x_lo, x_hi, x_frac)) in x_coords.iter().enumerate() {
2699 let tl = row_lo[x_lo];
2700 let tr = row_lo[x_hi];
2701 let bl = row_hi[x_lo];
2702 let br = row_hi[x_hi];
2703
2704 if (tl & tr & bl & br) < 0 {
2708 continue;
2710 }
2711 if tl > 0 && tr > 0 && bl > 0 && br > 0 {
2712 out_row[xi] = 255;
2714 continue;
2715 }
2716
2717 let x_frac_inv = FRAC_SCALE - x_frac;
2719 let l0 = tl as i64 * x_frac_inv as i64 + tr as i64 * x_frac as i64;
2720 let l1 = bl as i64 * x_frac_inv as i64 + br as i64 * x_frac as i64;
2721 let logit = l0 * y_frac_inv as i64 + l1 * y_frac as i64;
2722 out_row[xi] = if logit > 0 { 255 } else { 0 };
2723 }
2724 }
2725
2726 let tile = ndarray::Array3::from_shape_vec((bbox_h, bbox_w, 1), tile_buf)
2727 .expect("tile_buf length matches bbox_h * bbox_w");
2728 Ok(edgefirst_tensor::Segmentation {
2729 xmin,
2730 ymin,
2731 xmax,
2732 ymax,
2733 segmentation: crate::cpu::masks::mask_to_tensor(tile.view())?,
2734 })
2735 })
2736 .collect()
2737}
2738
2739#[allow(clippy::too_many_arguments)]
2740fn scaled_run<P: Copy + Sync>(
2741 detect: &[crate::DetectBox],
2742 coeff_all: &[f32],
2743 protos: &[P],
2744 proto_h: usize,
2745 proto_w: usize,
2746 num_protos: usize,
2747 letterbox: Option<[f32; 4]>,
2748 width: u32,
2749 height: u32,
2750 acc_scale: f32,
2751 load_f32: impl Fn(&P, f32) -> f32 + Copy + Sync,
2752) -> crate::Result<Vec<edgefirst_tensor::Segmentation>> {
2753 let (lx0, lw, ly0, lh) = match letterbox {
2754 Some([lx0, ly0, lx1, ly1]) => {
2755 let lw = if lx1 > lx0 { lx1 - lx0 } else { 1.0 };
2763 let lh = if ly1 > ly0 { ly1 - ly0 } else { 1.0 };
2764 (lx0, lw, ly0, lh)
2765 }
2766 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
2767 };
2768 let out_w = width as usize;
2769 let out_h = height as usize;
2770 let stride_y = proto_w * num_protos;
2771
2772 detect
2794 .par_iter()
2795 .enumerate()
2796 .map(|(i, det)| {
2797 let coeff = &coeff_all[i * num_protos..(i + 1) * num_protos];
2798 let bbox = det.bbox.to_canonical();
2799 let xmin = ((bbox.xmin - lx0) / lw).clamp(0.0, 1.0);
2800 let ymin = ((bbox.ymin - ly0) / lh).clamp(0.0, 1.0);
2801 let xmax = ((bbox.xmax - lx0) / lw).clamp(0.0, 1.0);
2802 let ymax = ((bbox.ymax - ly0) / lh).clamp(0.0, 1.0);
2803 let px0 = (xmin * out_w as f32).round() as usize;
2804 let py0 = (ymin * out_h as f32).round() as usize;
2805 let px1 = ((xmax * out_w as f32).round() as usize).min(out_w);
2806 let py1 = ((ymax * out_h as f32).round() as usize).min(out_h);
2807 let bbox_w = px1.saturating_sub(px0).max(1);
2808 let bbox_h = py1.saturating_sub(py0).max(1);
2809
2810 let sample_x_at = |px: f32| -> f32 {
2815 let model_x_norm = lx0 + (px + 0.5) / out_w as f32 * lw;
2816 model_x_norm * proto_w as f32 - 0.5
2817 };
2818 let sample_y_at = |py: f32| -> f32 {
2819 let model_y_norm = ly0 + (py + 0.5) / out_h as f32 * lh;
2820 model_y_norm * proto_h as f32 - 0.5
2821 };
2822 let s_x_min = sample_x_at(px0 as f32);
2823 let s_x_max = sample_x_at((px1 as f32) - 1.0);
2824 let s_y_min = sample_y_at(py0 as f32);
2825 let s_y_max = sample_y_at((py1 as f32) - 1.0);
2826 let proto_x0 = (s_x_min.floor() as isize)
2830 .max(0)
2831 .min(proto_w.saturating_sub(1) as isize) as usize;
2832 let proto_x1 = ((s_x_max.ceil() as isize) + 1).max(0).min(proto_w as isize) as usize;
2833 let proto_y0 = (s_y_min.floor() as isize)
2834 .max(0)
2835 .min(proto_h.saturating_sub(1) as isize) as usize;
2836 let proto_y1 = ((s_y_max.ceil() as isize) + 1).max(0).min(proto_h as isize) as usize;
2837 let roi_w = proto_x1.saturating_sub(proto_x0).max(1);
2838 let roi_h = proto_y1.saturating_sub(proto_y0).max(1);
2839
2840 if !acc_scale.is_finite() || acc_scale <= 0.0 {
2849 return Err(crate::Error::NotSupported(format!(
2850 "acc_scale must be finite and positive for sign-threshold optimization (got {acc_scale})"
2851 )));
2852 }
2853 let _ = acc_scale; let mut logits = vec![0.0_f32; roi_h * roi_w];
2855 for ly_idx in 0..roi_h {
2856 let py = proto_y0 + ly_idx;
2857 let row_base = py * stride_y + proto_x0 * num_protos;
2858 for lx_idx in 0..roi_w {
2859 let pix_base = row_base + lx_idx * num_protos;
2860 let mut acc = 0.0_f32;
2861 let mut k = 0;
2863 let chunks = num_protos / 4;
2864 for _ in 0..chunks {
2865 acc += coeff[k] * load_f32(&protos[pix_base + k], 0.0)
2866 + coeff[k + 1] * load_f32(&protos[pix_base + k + 1], 0.0)
2867 + coeff[k + 2] * load_f32(&protos[pix_base + k + 2], 0.0)
2868 + coeff[k + 3] * load_f32(&protos[pix_base + k + 3], 0.0);
2869 k += 4;
2870 }
2871 while k < num_protos {
2872 acc += coeff[k] * load_f32(&protos[pix_base + k], 0.0);
2873 k += 1;
2874 }
2875 logits[ly_idx * roi_w + lx_idx] = acc;
2876 }
2877 }
2878
2879 let roi_last_x = roi_w.saturating_sub(1);
2890 let roi_last_y = roi_h.saturating_sub(1);
2891
2892 let x_coords: Vec<(u32, u32, f32)> = (0..bbox_w)
2894 .map(|xi| {
2895 let sample_x = sample_x_at((px0 + xi) as f32) - proto_x0 as f32;
2896 let x_floor = sample_x.floor();
2897 let x_lo = (x_floor as isize).max(0).min(roi_last_x as isize) as u32;
2898 let x_hi = (x_lo as usize + 1).min(roi_w - 1) as u32;
2899 let x_frac = (sample_x - x_floor).clamp(0.0, 1.0);
2900 (x_lo, x_hi, x_frac)
2901 })
2902 .collect();
2903
2904 let mut tile_buf = vec![0u8; bbox_h * bbox_w];
2907 for yi in 0..bbox_h {
2908 let sample_y = sample_y_at((py0 + yi) as f32) - proto_y0 as f32;
2909 let y_floor = sample_y.floor();
2910 let y_lo = (y_floor as isize).max(0).min(roi_last_y as isize) as usize;
2911 let y_hi = (y_lo + 1).min(roi_h - 1);
2912 let y_frac = (sample_y - y_floor).clamp(0.0, 1.0);
2913 let row_lo = &logits[y_lo * roi_w..y_lo * roi_w + roi_w];
2914 let row_hi = &logits[y_hi * roi_w..y_hi * roi_w + roi_w];
2915 let out_row = &mut tile_buf[yi * bbox_w..(yi + 1) * bbox_w];
2916 for (xi, &(x_lo, x_hi, x_frac)) in x_coords.iter().enumerate() {
2917 let (xl, xh) = (x_lo as usize, x_hi as usize);
2918 let l0 = row_lo[xl] + (row_lo[xh] - row_lo[xl]) * x_frac;
2919 let l1 = row_hi[xl] + (row_hi[xh] - row_hi[xl]) * x_frac;
2920 let logit = l0 + (l1 - l0) * y_frac;
2921 out_row[xi] = if logit > 0.0 { 255 } else { 0 };
2922 }
2923 }
2924 let tile = ndarray::Array3::from_shape_vec((bbox_h, bbox_w, 1), tile_buf)
2926 .expect("tile_buf length matches bbox_h * bbox_w");
2927 Ok(edgefirst_tensor::Segmentation {
2928 xmin,
2929 ymin,
2930 xmax,
2931 ymax,
2932 segmentation: crate::cpu::masks::mask_to_tensor(tile.view())?,
2933 })
2934 })
2935 .collect()
2936}
2937
2938#[cfg(test)]
2939mod tests {
2940 use super::CPUProcessor;
2941 use edgefirst_tensor::{BoundingBox, DetectBox, ProtoData, ProtoLayout};
2942 use edgefirst_tensor::{Quantization, Tensor, TensorDyn};
2943
2944 const PROTO_H: usize = 4;
2945 const PROTO_W: usize = 4;
2946 const NUM_PROTOS: usize = 8;
2947
2948 fn det(xmin: f32, ymin: f32, xmax: f32, ymax: f32) -> DetectBox {
2949 DetectBox {
2950 bbox: BoundingBox {
2951 xmin,
2952 ymin,
2953 xmax,
2954 ymax,
2955 },
2956 score: 0.9,
2957 label: 0,
2958 }
2959 }
2960
2961 fn make_i8_quant(shape: &[usize], data: &[i8], scale: f32, zp: i32) -> TensorDyn {
2962 let t = Tensor::<i8>::from_slice(data, shape).unwrap();
2963 let t = t
2964 .with_quantization(Quantization::per_tensor(scale, zp))
2965 .unwrap();
2966 t.into()
2967 }
2968
2969 fn make_i16_quant(shape: &[usize], data: &[i16], scale: f32, zp: i32) -> TensorDyn {
2970 let t = Tensor::<i16>::from_slice(data, shape).unwrap();
2971 let t = t
2972 .with_quantization(Quantization::per_tensor(scale, zp))
2973 .unwrap();
2974 t.into()
2975 }
2976
2977 fn make_i16_raw(shape: &[usize], data: &[i16]) -> TensorDyn {
2978 let t = Tensor::<i16>::from_slice(data, shape).unwrap();
2979 t.into()
2980 }
2981
2982 fn make_f32(shape: &[usize], data: &[f32]) -> TensorDyn {
2983 let t = Tensor::<f32>::from_slice(data, shape).unwrap();
2984 t.into()
2985 }
2986
2987 fn gen_protos_i8(h: usize, w: usize, k: usize) -> Vec<i8> {
2988 (0..h * w * k).map(|i| (i % 127) as i8).collect()
2989 }
2990
2991 fn gen_coeffs_i16(n: usize, k: usize) -> Vec<i16> {
2992 (0..n * k)
2993 .map(|i| ((i as i32 % 201) - 100) as i16)
2994 .collect()
2995 }
2996
2997 fn gen_coeffs_i8(n: usize, k: usize) -> Vec<i8> {
2998 (0..n * k).map(|i| ((i as i32 % 201) - 100) as i8).collect()
2999 }
3000
3001 #[test]
3004 fn materialize_proto_i16_i8_quant_produces_masks() {
3005 let cpu = CPUProcessor::new();
3006 let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3007 let protos = make_i8_quant(
3008 &[PROTO_H, PROTO_W, NUM_PROTOS],
3009 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3010 0.02,
3011 0,
3012 );
3013 let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 0);
3014 let proto_data = ProtoData {
3015 mask_coefficients: coeffs,
3016 protos,
3017 layout: ProtoLayout::Nhwc,
3018 };
3019 let result = cpu.materialize_segmentations(&detect, &proto_data, None);
3020 assert!(result.is_ok(), "materialize failed: {:?}", result.err());
3021 let segs = result.unwrap();
3022 assert_eq!(segs.len(), 1);
3023 let seg = &segs[0];
3024 assert!(seg.segmentation.shape()[0] > 0);
3025 assert!(seg.segmentation.shape()[1] > 0);
3026 }
3027
3028 #[test]
3031 fn materialize_proto_i16_no_quant_falls_back_to_f32() {
3032 let cpu = CPUProcessor::new();
3033 let detect = vec![det(0.2, 0.2, 0.8, 0.8)];
3034 let protos = make_i8_quant(
3035 &[PROTO_H, PROTO_W, NUM_PROTOS],
3036 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3037 0.02,
3038 0,
3039 );
3040 let coeffs = make_i16_raw(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS));
3043 let proto_data = ProtoData {
3044 mask_coefficients: coeffs,
3045 protos,
3046 layout: ProtoLayout::Nhwc,
3047 };
3048 let result = cpu.materialize_segmentations(&detect, &proto_data, None);
3049 assert!(
3050 result.is_ok(),
3051 "missing coeff quant should fall back to f32 path, got: {:?}",
3052 result.err()
3053 );
3054 assert_eq!(result.unwrap().len(), 1);
3055 }
3056
3057 #[test]
3060 fn materialize_scaled_i16_i8_quant_produces_masks() {
3061 let cpu = CPUProcessor::new();
3062 let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3063 let protos = make_i8_quant(
3064 &[PROTO_H, PROTO_W, NUM_PROTOS],
3065 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3066 0.02,
3067 0,
3068 );
3069 let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 0);
3070 let proto_data = ProtoData {
3071 mask_coefficients: coeffs,
3072 protos,
3073 layout: ProtoLayout::Nhwc,
3074 };
3075 let result = cpu.materialize_scaled_segmentations(&detect, &proto_data, None, 64, 64);
3076 assert!(
3077 result.is_ok(),
3078 "materialize_scaled failed: {:?}",
3079 result.err()
3080 );
3081 let segs = result.unwrap();
3082 assert_eq!(segs.len(), 1);
3083 let seg = &segs[0];
3084 assert!(seg.segmentation.shape()[0] > 0);
3085 assert!(seg.segmentation.shape()[1] > 0);
3086 }
3087
3088 #[test]
3091 fn materialize_scaled_i16_no_quant_falls_back_to_f32() {
3092 let cpu = CPUProcessor::new();
3093 let detect = vec![det(0.2, 0.2, 0.8, 0.8)];
3094 let protos = make_i8_quant(
3095 &[PROTO_H, PROTO_W, NUM_PROTOS],
3096 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3097 0.02,
3098 0,
3099 );
3100 let coeffs = make_i16_raw(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS));
3101 let proto_data = ProtoData {
3102 mask_coefficients: coeffs,
3103 protos,
3104 layout: ProtoLayout::Nhwc,
3105 };
3106 let result = cpu.materialize_scaled_segmentations(&detect, &proto_data, None, 64, 64);
3107 assert!(
3108 result.is_ok(),
3109 "missing coeff quant should fall back to f32 path, got: {:?}",
3110 result.err()
3111 );
3112 assert_eq!(result.unwrap().len(), 1);
3113 }
3114
3115 #[test]
3118 fn materialize_proto_i16_i8_matches_f32_reference() {
3119 let cpu = CPUProcessor::new();
3120 let detect = vec![det(0.1, 0.1, 0.9, 0.9), det(0.3, 0.3, 0.7, 0.7)];
3121 let n_det = detect.len();
3122 let scale_c = 0.01_f32;
3123 let scale_p = 0.02_f32;
3124 let raw_protos = gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS);
3125 let raw_coeffs = gen_coeffs_i16(n_det, NUM_PROTOS);
3126
3127 let protos_f32: Vec<f32> = raw_protos.iter().map(|&v| v as f32 * scale_p).collect();
3129 let coeffs_f32: Vec<f32> = raw_coeffs.iter().map(|&v| v as f32 * scale_c).collect();
3130 let proto_data_f32 = ProtoData {
3131 mask_coefficients: make_f32(&[n_det, NUM_PROTOS], &coeffs_f32),
3132 protos: make_f32(&[PROTO_H, PROTO_W, NUM_PROTOS], &protos_f32),
3133 layout: ProtoLayout::Nhwc,
3134 };
3135
3136 let proto_data_int = ProtoData {
3137 mask_coefficients: make_i16_quant(&[n_det, NUM_PROTOS], &raw_coeffs, scale_c, 0),
3138 protos: make_i8_quant(&[PROTO_H, PROTO_W, NUM_PROTOS], &raw_protos, scale_p, 0),
3139 layout: ProtoLayout::Nhwc,
3140 };
3141
3142 let segs_f32 = cpu
3143 .materialize_segmentations(&detect, &proto_data_f32, None)
3144 .unwrap();
3145 let segs_int = cpu
3146 .materialize_segmentations(&detect, &proto_data_int, None)
3147 .unwrap();
3148
3149 assert_eq!(segs_f32.len(), segs_int.len());
3150 for (sf, si) in segs_f32.iter().zip(segs_int.iter()) {
3151 assert_eq!(sf.segmentation.shape(), si.segmentation.shape());
3152 use edgefirst_tensor::{TensorMapTrait as _, TensorTrait as _};
3153 let mf = sf
3154 .segmentation
3155 .as_typed::<u8>()
3156 .unwrap()
3157 .map_read()
3158 .unwrap();
3159 let mi = si
3160 .segmentation
3161 .as_typed::<u8>()
3162 .unwrap()
3163 .map_read()
3164 .unwrap();
3165 let (af, ai) = (mf.as_slice(), mi.as_slice());
3166 let total = af.len();
3167 let mismatches = af.iter().zip(ai.iter()).filter(|(a, b)| a != b).count();
3168 let pct = mismatches as f64 / total as f64 * 100.0;
3169 assert!(
3170 pct < 5.0,
3171 "mask mismatch {mismatches}/{total} ({pct:.1}%) exceeds 5% threshold"
3172 );
3173 }
3174 }
3175
3176 #[test]
3179 fn materialize_proto_i16_multiple_detections() {
3180 let cpu = CPUProcessor::new();
3181 let detect = vec![
3182 det(0.0, 0.0, 0.5, 0.5),
3183 det(0.5, 0.5, 1.0, 1.0),
3184 det(0.1, 0.1, 0.3, 0.3),
3185 ];
3186 let protos = make_i8_quant(
3187 &[PROTO_H, PROTO_W, NUM_PROTOS],
3188 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3189 0.02,
3190 0,
3191 );
3192 let coeffs = make_i16_quant(&[3, NUM_PROTOS], &gen_coeffs_i16(3, NUM_PROTOS), 0.01, 0);
3193 let proto_data = ProtoData {
3194 mask_coefficients: coeffs,
3195 protos,
3196 layout: ProtoLayout::Nhwc,
3197 };
3198 let segs = cpu
3199 .materialize_segmentations(&detect, &proto_data, None)
3200 .unwrap();
3201 assert_eq!(segs.len(), 3);
3202 }
3203
3204 #[test]
3207 fn materialize_proto_i16_empty_detections() {
3208 let cpu = CPUProcessor::new();
3209 let detect: Vec<DetectBox> = vec![];
3210 let protos = make_i8_quant(
3211 &[PROTO_H, PROTO_W, NUM_PROTOS],
3212 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3213 0.02,
3214 0,
3215 );
3216 let coeffs = make_i16_quant(&[0, NUM_PROTOS], &[], 0.01, 0);
3217 let proto_data = ProtoData {
3218 mask_coefficients: coeffs,
3219 protos,
3220 layout: ProtoLayout::Nhwc,
3221 };
3222 let segs = cpu
3223 .materialize_segmentations(&detect, &proto_data, None)
3224 .unwrap();
3225 assert!(segs.is_empty());
3226 }
3227
3228 #[test]
3231 fn materialize_scaled_i16_i8_matches_f32_reference() {
3232 let cpu = CPUProcessor::new();
3233 let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3234 let scale_c = 0.01_f32;
3235 let scale_p = 0.02_f32;
3236 let raw_protos = gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS);
3237 let raw_coeffs = gen_coeffs_i16(1, NUM_PROTOS);
3238
3239 let protos_f32: Vec<f32> = raw_protos.iter().map(|&v| v as f32 * scale_p).collect();
3240 let coeffs_f32: Vec<f32> = raw_coeffs.iter().map(|&v| v as f32 * scale_c).collect();
3241 let proto_data_f32 = ProtoData {
3242 mask_coefficients: make_f32(&[1, NUM_PROTOS], &coeffs_f32),
3243 protos: make_f32(&[PROTO_H, PROTO_W, NUM_PROTOS], &protos_f32),
3244 layout: ProtoLayout::Nhwc,
3245 };
3246 let proto_data_int = ProtoData {
3247 mask_coefficients: make_i16_quant(&[1, NUM_PROTOS], &raw_coeffs, scale_c, 0),
3248 protos: make_i8_quant(&[PROTO_H, PROTO_W, NUM_PROTOS], &raw_protos, scale_p, 0),
3249 layout: ProtoLayout::Nhwc,
3250 };
3251
3252 let (w, h) = (64_u32, 64_u32);
3253 let segs_f32 = cpu
3254 .materialize_scaled_segmentations(&detect, &proto_data_f32, None, w, h)
3255 .unwrap();
3256 let segs_int = cpu
3257 .materialize_scaled_segmentations(&detect, &proto_data_int, None, w, h)
3258 .unwrap();
3259
3260 assert_eq!(segs_f32.len(), segs_int.len());
3261 for (sf, si) in segs_f32.iter().zip(segs_int.iter()) {
3262 assert_eq!(sf.segmentation.shape(), si.segmentation.shape());
3263 use edgefirst_tensor::{TensorMapTrait as _, TensorTrait as _};
3264 let mf = sf
3265 .segmentation
3266 .as_typed::<u8>()
3267 .unwrap()
3268 .map_read()
3269 .unwrap();
3270 let mi = si
3271 .segmentation
3272 .as_typed::<u8>()
3273 .unwrap()
3274 .map_read()
3275 .unwrap();
3276 let (af, ai) = (mf.as_slice(), mi.as_slice());
3277 let total = af.len();
3278 let mismatches = af.iter().zip(ai.iter()).filter(|(a, b)| a != b).count();
3279 let pct = mismatches as f64 / total as f64 * 100.0;
3280 assert!(
3281 pct < 5.0,
3282 "scaled mask mismatch {mismatches}/{total} ({pct:.1}%) exceeds 5% threshold"
3283 );
3284 }
3285 }
3286
3287 #[test]
3290 fn materialize_proto_i8_i8_regression() {
3291 let cpu = CPUProcessor::new();
3292 let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3293 let protos = make_i8_quant(
3294 &[PROTO_H, PROTO_W, NUM_PROTOS],
3295 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3296 0.02,
3297 0,
3298 );
3299 let coeffs = make_i8_quant(&[1, NUM_PROTOS], &gen_coeffs_i8(1, NUM_PROTOS), 0.01, 0);
3300 let proto_data = ProtoData {
3301 mask_coefficients: coeffs,
3302 protos,
3303 layout: ProtoLayout::Nhwc,
3304 };
3305 let result = cpu.materialize_segmentations(&detect, &proto_data, None);
3306 assert!(result.is_ok(), "i8×i8 regression: {:?}", result.err());
3307 assert_eq!(result.unwrap().len(), 1);
3308 }
3309
3310 #[test]
3313 fn materialize_proto_i16_nonzero_zp() {
3314 let cpu = CPUProcessor::new();
3315 let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3316 let protos = make_i8_quant(
3317 &[PROTO_H, PROTO_W, NUM_PROTOS],
3318 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3319 0.02,
3320 -10,
3321 );
3322 let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 5);
3323 let proto_data = ProtoData {
3324 mask_coefficients: coeffs,
3325 protos,
3326 layout: ProtoLayout::Nhwc,
3327 };
3328 let result = cpu.materialize_segmentations(&detect, &proto_data, None);
3329 assert!(result.is_ok(), "nonzero zp failed: {:?}", result.err());
3330 assert_eq!(result.unwrap().len(), 1);
3331 }
3332
3333 #[test]
3336 fn materialize_scaled_i16_nonzero_zp() {
3337 let cpu = CPUProcessor::new();
3338 let detect = vec![det(0.1, 0.1, 0.9, 0.9)];
3339 let protos = make_i8_quant(
3340 &[PROTO_H, PROTO_W, NUM_PROTOS],
3341 &gen_protos_i8(PROTO_H, PROTO_W, NUM_PROTOS),
3342 0.02,
3343 -10,
3344 );
3345 let coeffs = make_i16_quant(&[1, NUM_PROTOS], &gen_coeffs_i16(1, NUM_PROTOS), 0.01, 5);
3346 let proto_data = ProtoData {
3347 mask_coefficients: coeffs,
3348 protos,
3349 layout: ProtoLayout::Nhwc,
3350 };
3351 let result = cpu.materialize_scaled_segmentations(&detect, &proto_data, None, 64, 64);
3352 assert!(
3353 result.is_ok(),
3354 "scaled nonzero zp failed: {:?}",
3355 result.err()
3356 );
3357 assert_eq!(result.unwrap().len(), 1);
3358 }
3359}