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