1use crate::internal::*;
2use num_traits::AsPrimitive;
3use std::iter::Sum;
4
5use crate::ops::cnn::pools::{ConcretePoolGeometry, PoolGeometry, PoolSpec};
6
7crate::declare_knob!(
8 TRACT_AVGPOOL_SEPARABLE,
9 bool,
10 false,
11 "Use the separable average-pool kernel for stride-1 NCHW/NHWC pools. Not bit-identical: \
12 it reassociates the sum, permitted by SumPool's Validation::Rounding contract."
13);
14
15#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
16pub struct SumPool {
17 pub pool_spec: PoolSpec,
18 pub count_include_pad: bool,
19 pub normalize: bool,
20}
21
22impl Op for SumPool {
23 fn name(&self) -> StaticName {
24 "SumPool".into()
25 }
26
27 fn info(&self) -> TractResult<Vec<String>> {
28 Ok(self.pool_spec.info())
29 }
30
31 fn validation(&self) -> Validation {
32 Validation::Rounding
33 }
34
35 op_as_typed_op!();
36}
37
38impl EvalOp for SumPool {
39 fn is_stateless(&self) -> bool {
40 true
41 }
42
43 fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
44 let shape: TVec<TDim> = inputs[0].shape().iter().map(|d| d.to_dim()).collect();
45 self.to_optimized(&shape)?.eval(inputs)
46 }
47}
48
49impl TypedOp for SumPool {
50 fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
51 self.pool_spec.output_facts(inputs)
52 }
53
54 fn declutter(
55 &self,
56 model: &TypedModel,
57 node: &TypedNode,
58 ) -> TractResult<Option<TypedModelPatch>> {
59 let fact = model.outlet_fact(node.inputs[0])?;
60 if let Some(pool_spec) = self.pool_spec.declutter(&fact.shape)? {
61 return Ok(Some(TypedModelPatch::replace_single_op(
62 model,
63 node,
64 &node.inputs,
65 Self { pool_spec, ..self.clone() },
66 )?));
67 }
68 Ok(None)
69 }
70
71 fn codegen(
75 &self,
76 model: &TypedModel,
77 node: &TypedNode,
78 ) -> TractResult<Option<TypedModelPatch>> {
79 let fact = model.outlet_fact(node.inputs[0])?;
80 if fact.shape.as_concrete().is_none() {
81 return Ok(None);
82 }
83 let mut op = self.to_optimized(&fact.shape.to_tvec())?;
84 op.geometry = op.geometry.optimize_if(fact.shape.as_concrete())?;
85 Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs, op)?))
86 }
87
88 as_op!();
89}
90
91impl SumPool {
92 fn to_optimized(&self, input_shape: &[TDim]) -> TractResult<OptSumPool> {
93 Ok(OptSumPool {
94 pool_spec: self.pool_spec.clone(),
95 count_include_pad: self.count_include_pad,
96 normalize: self.normalize,
97 geometry: self.pool_spec.compute_geo(input_shape)?,
98 })
99 }
100}
101
102#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
103pub struct OptSumPool {
104 pub pool_spec: PoolSpec,
105 pub count_include_pad: bool,
106 pub normalize: bool,
107 pub geometry: PoolGeometry,
108}
109
110impl Op for OptSumPool {
111 fn name(&self) -> StaticName {
112 "OptSumPool".into()
113 }
114
115 fn info(&self) -> TractResult<Vec<String>> {
116 Ok(self.pool_spec.info())
117 }
118
119 fn validation(&self) -> Validation {
120 Validation::Rounding
121 }
122
123 op_as_typed_op!();
124}
125
126impl EvalOp for OptSumPool {
127 fn is_stateless(&self) -> bool {
128 true
129 }
130
131 fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
132 let input = args_1!(inputs);
133 let geo = self.geometry.to_concrete(input.shape())?;
134 let values = if input.datum_type().is_float() {
135 let mut values =
136 unsafe { Tensor::uninitialized_dt(input.datum_type(), &geo.output_shape.shape)? };
137 dispatch_floatlike!(Self::eval_t(input.datum_type())(
138 self,
139 &*input,
140 values.as_ptr_mut()?,
141 geo.as_ref()
142 ))?;
143 values
144 } else {
145 let mut values =
146 unsafe { Tensor::uninitialized_dt(DatumType::F32, &geo.output_shape.shape)? };
147 let input_f32 = input.cast_to_dt(DatumType::F32)?;
148 self.eval_t::<f32>(input_f32.as_ref(), values.as_ptr_mut()?, geo.as_ref())?;
149 values.cast_to_dt(input.datum_type())?.into_owned()
150 };
151
152 Ok(tvec!(values.into_tvalue()))
153 }
154}
155
156impl TypedOp for OptSumPool {
157 fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
158 self.pool_spec.output_facts(inputs)
159 }
160
161 fn declutter(
162 &self,
163 model: &TypedModel,
164 node: &TypedNode,
165 ) -> TractResult<Option<TypedModelPatch>> {
166 let fact = model.outlet_fact(node.inputs[0])?;
167 if let Some(pool_spec) = self.pool_spec.declutter(&fact.shape)? {
168 return Ok(Some(TypedModelPatch::replace_single_op(
169 model,
170 node,
171 &node.inputs,
172 Self { pool_spec, ..self.clone() },
173 )?));
174 }
175 Ok(None)
176 }
177
178 as_op!();
179}
180
181impl OptSumPool {
182 fn eval_t<T: Copy + Datum + Sum + num_traits::Float>(
183 &self,
184 input: &Tensor,
185 values_ptr: *mut T,
186 geo: &ConcretePoolGeometry,
187 ) -> TractResult<()>
188 where
189 usize: AsPrimitive<T>,
190 {
191 if self.try_fast_2d::<T>(input, values_ptr, geo)? {
192 return Ok(());
193 }
194 let input_ptr = input.as_ptr::<T>()?;
195
196 let n = *geo.input_shape.n().unwrap_or(&1);
197 let n_stride_i = geo.input_shape.n_stride().unwrap_or(&0);
198 let n_stride_o = geo.output_shape.n_stride().unwrap_or(&0);
199 unsafe {
200 geo.patch.visit_output(|visitor| {
201 let div: Option<T> = if self.normalize {
202 Some(
203 if self.count_include_pad {
204 geo.patch.standard_layout_data_field.len().as_()
205 } else {
206 visitor.valid_count().as_()
207 }
208 .recip(),
209 )
210 } else {
211 None
212 };
213 for n in 0..n {
214 let input_offset = n * n_stride_i;
215 let output_offset = n * n_stride_o;
216 for c in 0..*geo.input_shape.c() {
217 let input_offset = input_offset + geo.input_shape.c_stride() * c;
218 let output_offset = output_offset + geo.output_shape.c_stride() * c;
219 let sum = visitor
220 .valid_offsets()
221 .map(|v| *input_ptr.offset(v + input_offset as isize))
222 .sum::<T>();
223
224 if let Some(div) = div {
225 *values_ptr.offset(output_offset as isize + visitor.output_offset) =
226 sum * div;
227 }
228 }
229 }
230 });
231 }
232 Ok(())
233 }
234
235 fn try_fast_2d<T: Copy + Datum + num_traits::Float>(
240 &self,
241 input: &Tensor,
242 values_ptr: *mut T,
243 geo: &ConcretePoolGeometry,
244 ) -> TractResult<bool>
245 where
246 usize: AsPrimitive<T>,
247 {
248 let patch = &geo.patch;
249 if !TRACT_AVGPOOL_SEPARABLE.get()
250 || !self.normalize
251 || patch.rank() != 2
252 || *patch.spec.strides != [1, 1]
253 || *patch.spec.dilations != [1, 1]
254 {
255 return Ok(false);
256 }
257 let input_ptr = input.as_ptr::<T>()?;
258 let ish = &geo.input_shape;
259 if *ish.w_stride() == 1 {
260 unsafe {
261 self.fast_2d_separable::<T>(input_ptr, values_ptr, geo);
262 }
263 Ok(true)
264 } else if *ish.c_stride() == 1 && *ish.w_stride() == *ish.c() {
265 unsafe {
266 self.fast_2d_separable_nhwc::<T>(input_ptr, values_ptr, geo);
267 }
268 Ok(true)
269 } else {
270 Ok(false)
271 }
272 }
273
274 unsafe fn fast_2d_separable<T: Copy + Datum + num_traits::Float>(
279 &self,
280 input_ptr: *const T,
281 values_ptr: *mut T,
282 geo: &ConcretePoolGeometry,
283 ) where
284 usize: AsPrimitive<T>,
285 {
286 let ish = &geo.input_shape;
287 let osh = &geo.output_shape;
288 let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
289 let (ho, wo) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
290 let (kh, kw) =
291 (geo.patch.spec.kernel_shape[0] as isize, geo.patch.spec.kernel_shape[1] as isize);
292 let (pt, pl) = (geo.patch.pad_before[0] as isize, geo.patch.pad_before[1] as isize);
293 let ih_stride = *ish.h_stride() as isize;
294 let oh_stride = *osh.h_stride() as isize;
295 let ow_stride = *osh.w_stride() as isize;
296 let n = *ish.n().unwrap_or(&1);
297 let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
298 let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
299 let c = *ish.c();
300 let ic_stride = *ish.c_stride() as isize;
301 let oc_stride = *osh.c_stride() as isize;
302
303 let axis_valid = |out: usize, k: isize, pad: isize, lim: isize| -> Vec<usize> {
304 (0..out)
305 .map(|o| {
306 let lo = o as isize - pad;
307 let start = (-lo).max(0);
308 let end = (lim - lo).min(k);
309 (end - start).max(0) as usize
310 })
311 .collect()
312 };
313 let kx_valid = axis_valid(wo, kw, pl, w);
314 let ky_valid = axis_valid(ho, kh, pt, h);
315 let full_recip: T = ((kh * kw) as usize).as_().recip();
316
317 let mut htmp = vec![T::zero(); h as usize * wo];
318 unsafe {
319 for nn in 0..n as isize {
320 for cc in 0..c as isize {
321 let in_base = nn * in_stride + cc * ic_stride;
322 let out_base = nn * on_stride + cc * oc_stride;
323 for y in 0..h {
324 let row = in_base + y * ih_stride;
325 let dst = y as usize * wo;
326 let mut acc = T::zero();
327 for kx in 0..kw {
328 let ix = -pl + kx;
329 if ix >= 0 && ix < w {
330 acc = acc + *input_ptr.offset(row + ix);
331 }
332 }
333 *htmp.get_unchecked_mut(dst) = acc;
334 for ox in 1..wo as isize {
335 let entering = ox - pl + kw - 1;
336 let leaving = ox - pl - 1;
337 if entering >= 0 && entering < w {
338 acc = acc + *input_ptr.offset(row + entering);
339 }
340 if leaving >= 0 && leaving < w {
341 acc = acc - *input_ptr.offset(row + leaving);
342 }
343 *htmp.get_unchecked_mut(dst + ox as usize) = acc;
344 }
345 }
346 for ox in 0..wo {
347 let mut acc = T::zero();
348 for ky in 0..kh {
349 let iy = -pt + ky;
350 if iy >= 0 && iy < h {
351 acc = acc + *htmp.get_unchecked(iy as usize * wo + ox);
352 }
353 }
354 let store = |oy: usize, acc: T| {
355 let div = if self.count_include_pad {
356 full_recip
357 } else {
358 (kx_valid[ox] * ky_valid[oy]).as_().recip()
359 };
360 *values_ptr.offset(
361 out_base + oy as isize * oh_stride + ox as isize * ow_stride,
362 ) = acc * div;
363 };
364 store(0, acc);
365 for oy in 1..ho as isize {
366 let entering = oy - pt + kh - 1;
367 let leaving = oy - pt - 1;
368 if entering >= 0 && entering < h {
369 acc = acc + *htmp.get_unchecked(entering as usize * wo + ox);
370 }
371 if leaving >= 0 && leaving < h {
372 acc = acc - *htmp.get_unchecked(leaving as usize * wo + ox);
373 }
374 store(oy as usize, acc);
375 }
376 }
377 }
378 }
379 }
380 }
381
382 unsafe fn fast_2d_separable_nhwc<T: Copy + Datum + num_traits::Float>(
387 &self,
388 input_ptr: *const T,
389 values_ptr: *mut T,
390 geo: &ConcretePoolGeometry,
391 ) where
392 usize: AsPrimitive<T>,
393 {
394 let ish = &geo.input_shape;
395 let osh = &geo.output_shape;
396 let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
397 let (ho, wo) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
398 let (kh, kw) =
399 (geo.patch.spec.kernel_shape[0] as isize, geo.patch.spec.kernel_shape[1] as isize);
400 let (pt, pl) = (geo.patch.pad_before[0] as isize, geo.patch.pad_before[1] as isize);
401 let ih_stride = *ish.h_stride() as isize;
402 let iw_stride = *ish.w_stride() as isize;
403 let oh_stride = *osh.h_stride() as isize;
404 let ow_stride = *osh.w_stride() as isize;
405 let n = *ish.n().unwrap_or(&1);
406 let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
407 let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
408 let c = *ish.c();
409
410 let axis_valid = |out: usize, k: isize, pad: isize, lim: isize| -> Vec<usize> {
411 (0..out)
412 .map(|o| {
413 let lo = o as isize - pad;
414 let start = (-lo).max(0);
415 let end = (lim - lo).min(k);
416 (end - start).max(0) as usize
417 })
418 .collect()
419 };
420 let kx_valid = axis_valid(wo, kw, pl, w);
421 let ky_valid = axis_valid(ho, kh, pt, h);
422 let full_recip: T = ((kh * kw) as usize).as_().recip();
423
424 let mut htmp = vec![T::zero(); h as usize * wo * c];
425 let mut acc = vec![T::zero(); c];
426 unsafe {
427 for nn in 0..n as isize {
428 let in_base = nn * in_stride;
429 let out_base = nn * on_stride;
430 for y in 0..h {
431 let row = in_base + y * ih_stride;
432 let hrow = y as usize * wo * c;
433 acc.iter_mut().for_each(|a| *a = T::zero());
434 for kx in 0..kw {
435 let ix = -pl + kx;
436 if ix >= 0 && ix < w {
437 let p = row + ix * iw_stride;
438 for (ch, a) in acc.iter_mut().enumerate() {
439 *a = *a + *input_ptr.offset(p + ch as isize);
440 }
441 }
442 }
443 htmp[hrow..hrow + c].copy_from_slice(&acc);
444 for ox in 1..wo as isize {
445 let entering = ox - pl + kw - 1;
446 let leaving = ox - pl - 1;
447 if entering >= 0 && entering < w {
448 let p = row + entering * iw_stride;
449 for (ch, a) in acc.iter_mut().enumerate() {
450 *a = *a + *input_ptr.offset(p + ch as isize);
451 }
452 }
453 if leaving >= 0 && leaving < w {
454 let p = row + leaving * iw_stride;
455 for (ch, a) in acc.iter_mut().enumerate() {
456 *a = *a - *input_ptr.offset(p + ch as isize);
457 }
458 }
459 let dst = hrow + ox as usize * c;
460 htmp[dst..dst + c].copy_from_slice(&acc);
461 }
462 }
463 for ox in 0..wo {
464 acc.iter_mut().for_each(|a| *a = T::zero());
465 for ky in 0..kh {
466 let iy = -pt + ky;
467 if iy >= 0 && iy < h {
468 let src = iy as usize * wo * c + ox * c;
469 for (ch, a) in acc.iter_mut().enumerate() {
470 *a = *a + *htmp.get_unchecked(src + ch);
471 }
472 }
473 }
474 let store = |oy: usize, acc: &[T]| {
475 let div = if self.count_include_pad {
476 full_recip
477 } else {
478 (kx_valid[ox] * ky_valid[oy]).as_().recip()
479 };
480 let o = out_base + oy as isize * oh_stride + ox as isize * ow_stride;
481 for (ch, &a) in acc.iter().enumerate() {
482 *values_ptr.offset(o + ch as isize) = a * div;
483 }
484 };
485 store(0, &acc);
486 for oy in 1..ho as isize {
487 let entering = oy - pt + kh - 1;
488 let leaving = oy - pt - 1;
489 if entering >= 0 && entering < h {
490 let src = entering as usize * wo * c + ox * c;
491 for (ch, a) in acc.iter_mut().enumerate() {
492 *a = *a + *htmp.get_unchecked(src + ch);
493 }
494 }
495 if leaving >= 0 && leaving < h {
496 let src = leaving as usize * wo * c + ox * c;
497 for (ch, a) in acc.iter_mut().enumerate() {
498 *a = *a - *htmp.get_unchecked(src + ch);
499 }
500 }
501 store(oy as usize, &acc);
502 }
503 }
504 }
505 }
506 }
507}
508
509#[cfg(test)]
510mod tests {
511 use super::*;
512 use crate::ops::cnn::PaddingSpec;
513 use crate::ops::nn::DataFormat;
514
515 fn test_case() -> (TypedModel, TVec<TValue>) {
516 let mut model = TypedModel::default();
517 let source = model.add_source("data", f32::fact([1, 3, 8, 8])).unwrap();
518 let pool_spec = PoolSpec::new(
519 DataFormat::NCHW,
520 tvec![2, 2],
521 PaddingSpec::Valid,
522 None,
523 Some(tvec![2, 2]),
524 3,
525 3,
526 );
527 let op = SumPool { pool_spec, count_include_pad: false, normalize: true };
528 let out = model.wire_node("pool", op, &[source]).unwrap();
529 model.select_output_outlets(&out).unwrap();
530 let input = ndarray::Array4::from_shape_fn((1, 3, 8, 8), |(_, c, y, x)| {
531 (c * 64 + y * 8 + x) as f32
532 })
533 .into_tensor()
534 .into_tvalue();
535 (model, tvec!(input))
536 }
537
538 #[test]
539 fn optimized_sumpool_has_concrete_geometry() {
540 let (model, input) = test_case();
541 let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
542
543 let optimized = model.into_optimized().unwrap();
544 let pool = optimized
545 .nodes
546 .iter()
547 .find_map(|n| n.op_as::<OptSumPool>())
548 .expect("optimized model should contain an OptSumPool");
549 assert!(
550 pool.geometry.is_concrete(),
551 "OptSumPool geometry should be concrete after optimization"
552 );
553
554 let opt = optimized.into_runnable().unwrap().run(input).unwrap();
555 assert_eq!(*opt[0], *plain[0]);
556 }
557
558 #[test]
559 fn separable_matches_generic_kernel() {
560 let (c, h, w) = (5usize, 7usize, 9usize);
561 let pool_spec = PoolSpec::new(
562 DataFormat::NCHW,
563 tvec![3, 3],
564 PaddingSpec::SameUpper,
565 None,
566 Some(tvec![1, 1]),
567 c,
568 c,
569 );
570 let op = OptSumPool {
571 pool_spec: pool_spec.clone(),
572 count_include_pad: false,
573 normalize: true,
574 geometry: pool_spec
575 .compute_geo(&[1.to_dim(), c.to_dim(), h.to_dim(), w.to_dim()])
576 .unwrap(),
577 };
578 let input: Tensor = ndarray::Array4::from_shape_fn((1, c, h, w), |(_, cc, y, x)| {
579 ((cc * 17 + y * 3 + x) % 13) as f32 - 6.0
580 })
581 .into_tensor();
582
583 let generic = op.eval(tvec![input.clone().into_tvalue()]).unwrap();
585 let generic = generic[0].try_as_plain().unwrap().as_slice::<f32>().unwrap().to_vec();
586
587 let geo = op.geometry.to_concrete(input.shape()).unwrap();
589 let mut out = Tensor::zero::<f32>(&geo.output_shape.shape).unwrap();
590 unsafe {
591 op.fast_2d_separable::<f32>(
592 input.as_ptr::<f32>().unwrap(),
593 out.as_ptr_mut::<f32>().unwrap(),
594 geo.as_ref(),
595 );
596 }
597 let sep = out.try_as_plain().unwrap().as_slice::<f32>().unwrap();
598
599 let max_abs = generic.iter().zip(sep).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
600 assert!(max_abs < 1e-4, "separable vs generic max abs diff {max_abs}");
601 }
602
603 #[test]
604 fn separable_nhwc_matches_generic_kernel() {
605 let (c, h, w) = (5usize, 7usize, 9usize);
606 let pool_spec = PoolSpec::new(
607 DataFormat::NHWC,
608 tvec![3, 3],
609 PaddingSpec::SameUpper,
610 None,
611 Some(tvec![1, 1]),
612 c,
613 c,
614 );
615 let op = OptSumPool {
616 pool_spec: pool_spec.clone(),
617 count_include_pad: false,
618 normalize: true,
619 geometry: pool_spec
620 .compute_geo(&[1.to_dim(), h.to_dim(), w.to_dim(), c.to_dim()])
621 .unwrap(),
622 };
623 let input: Tensor = ndarray::Array4::from_shape_fn((1, h, w, c), |(_, y, x, cc)| {
624 ((cc * 17 + y * 3 + x) % 13) as f32 - 6.0
625 })
626 .into_tensor();
627
628 let generic = op.eval(tvec![input.clone().into_tvalue()]).unwrap();
630 let generic = generic[0].try_as_plain().unwrap().as_slice::<f32>().unwrap().to_vec();
631
632 let geo = op.geometry.to_concrete(input.shape()).unwrap();
634 let mut out = Tensor::zero::<f32>(&geo.output_shape.shape).unwrap();
635 unsafe {
636 op.fast_2d_separable_nhwc::<f32>(
637 input.as_ptr::<f32>().unwrap(),
638 out.as_ptr_mut::<f32>().unwrap(),
639 geo.as_ref(),
640 );
641 }
642 let sep = out.try_as_plain().unwrap().as_slice::<f32>().unwrap();
643
644 let max_abs = generic.iter().zip(sep).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
645 assert!(max_abs < 1e-4, "separable NHWC vs generic max abs diff {max_abs}");
646 }
647}