tract_core/ops/cnn/
sumpool.rs

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    /// Lower to `OptSumPool` with the geometry pre-resolved to `Concrete` when the
72    /// input shape is fixed, so the `Patch` is built once here rather than per eval.
73    /// Symbolic shapes are left as `SumPool`.
74    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                        *values_ptr.offset(output_offset as isize + visitor.output_offset) =
225                            if let Some(div) = div { sum * div } else { sum };
226                    }
227                }
228            });
229        }
230        Ok(())
231    }
232
233    /// Opt-in separable average-pool fast path, gated by `TRACT_AVGPOOL_SEPARABLE`.
234    /// Returns `true` when it handled the eval, `false` to fall back to the generic
235    /// kernel. Restricted to rank-2, stride-1, dilation-1, normalize pools; the NCHW
236    /// and NHWC layouts have dedicated kernels, any other layout falls back.
237    fn try_fast_2d<T: Copy + Datum + num_traits::Float>(
238        &self,
239        input: &Tensor,
240        values_ptr: *mut T,
241        geo: &ConcretePoolGeometry,
242    ) -> TractResult<bool>
243    where
244        usize: AsPrimitive<T>,
245    {
246        let patch = &geo.patch;
247        if !TRACT_AVGPOOL_SEPARABLE.get()
248            || !self.normalize
249            || patch.rank() != 2
250            || *patch.spec.strides != [1, 1]
251            || *patch.spec.dilations != [1, 1]
252        {
253            return Ok(false);
254        }
255        let input_ptr = input.as_ptr::<T>()?;
256        let ish = &geo.input_shape;
257        if *ish.w_stride() == 1 {
258            unsafe {
259                self.fast_2d_separable::<T>(input_ptr, values_ptr, geo);
260            }
261            Ok(true)
262        } else if *ish.c_stride() == 1 && *ish.w_stride() == *ish.c() {
263            unsafe {
264                self.fast_2d_separable_nhwc::<T>(input_ptr, values_ptr, geo);
265            }
266            Ok(true)
267        } else {
268            Ok(false)
269        }
270    }
271
272    /// Separable running-sum average pool. Out-of-bounds taps contribute 0, so the
273    /// box sum over the padded window equals the sum of in-bounds values; the divisor
274    /// is the per-cell valid count, itself separable into `kx_valid * ky_valid`.
275    /// Reassociates the sum, so it is not bit-identical to the generic kernel.
276    unsafe fn fast_2d_separable<T: Copy + Datum + num_traits::Float>(
277        &self,
278        input_ptr: *const T,
279        values_ptr: *mut T,
280        geo: &ConcretePoolGeometry,
281    ) where
282        usize: AsPrimitive<T>,
283    {
284        let ish = &geo.input_shape;
285        let osh = &geo.output_shape;
286        let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
287        let (ho, wo) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
288        let (kh, kw) =
289            (geo.patch.spec.kernel_shape[0] as isize, geo.patch.spec.kernel_shape[1] as isize);
290        let (pt, pl) = (geo.patch.pad_before[0] as isize, geo.patch.pad_before[1] as isize);
291        let ih_stride = *ish.h_stride() as isize;
292        let oh_stride = *osh.h_stride() as isize;
293        let ow_stride = *osh.w_stride() as isize;
294        let n = *ish.n().unwrap_or(&1);
295        let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
296        let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
297        let c = *ish.c();
298        let ic_stride = *ish.c_stride() as isize;
299        let oc_stride = *osh.c_stride() as isize;
300
301        let axis_valid = |out: usize, k: isize, pad: isize, lim: isize| -> Vec<usize> {
302            (0..out)
303                .map(|o| {
304                    let lo = o as isize - pad;
305                    let start = (-lo).max(0);
306                    let end = (lim - lo).min(k);
307                    (end - start).max(0) as usize
308                })
309                .collect()
310        };
311        let kx_valid = axis_valid(wo, kw, pl, w);
312        let ky_valid = axis_valid(ho, kh, pt, h);
313        let full_recip: T = ((kh * kw) as usize).as_().recip();
314
315        let mut htmp = vec![T::zero(); h as usize * wo];
316        unsafe {
317            for nn in 0..n as isize {
318                for cc in 0..c as isize {
319                    let in_base = nn * in_stride + cc * ic_stride;
320                    let out_base = nn * on_stride + cc * oc_stride;
321                    for y in 0..h {
322                        let row = in_base + y * ih_stride;
323                        let dst = y as usize * wo;
324                        let mut acc = T::zero();
325                        for kx in 0..kw {
326                            let ix = -pl + kx;
327                            if ix >= 0 && ix < w {
328                                acc = acc + *input_ptr.offset(row + ix);
329                            }
330                        }
331                        *htmp.get_unchecked_mut(dst) = acc;
332                        for ox in 1..wo as isize {
333                            let entering = ox - pl + kw - 1;
334                            let leaving = ox - pl - 1;
335                            if entering >= 0 && entering < w {
336                                acc = acc + *input_ptr.offset(row + entering);
337                            }
338                            if leaving >= 0 && leaving < w {
339                                acc = acc - *input_ptr.offset(row + leaving);
340                            }
341                            *htmp.get_unchecked_mut(dst + ox as usize) = acc;
342                        }
343                    }
344                    #[allow(clippy::needless_range_loop)]
345                    for ox in 0..wo {
346                        let mut acc = T::zero();
347                        for ky in 0..kh {
348                            let iy = -pt + ky;
349                            if iy >= 0 && iy < h {
350                                acc = acc + *htmp.get_unchecked(iy as usize * wo + ox);
351                            }
352                        }
353                        let store = |oy: usize, acc: T| {
354                            let div = if self.count_include_pad {
355                                full_recip
356                            } else {
357                                (kx_valid[ox] * ky_valid[oy]).as_().recip()
358                            };
359                            *values_ptr.offset(
360                                out_base + oy as isize * oh_stride + ox as isize * ow_stride,
361                            ) = acc * div;
362                        };
363                        store(0, acc);
364                        for oy in 1..ho as isize {
365                            let entering = oy - pt + kh - 1;
366                            let leaving = oy - pt - 1;
367                            if entering >= 0 && entering < h {
368                                acc = acc + *htmp.get_unchecked(entering as usize * wo + ox);
369                            }
370                            if leaving >= 0 && leaving < h {
371                                acc = acc - *htmp.get_unchecked(leaving as usize * wo + ox);
372                            }
373                            store(oy as usize, acc);
374                        }
375                    }
376                }
377            }
378        }
379    }
380
381    /// NHWC counterpart of `fast_2d_separable`. Channels are the innermost
382    /// (contiguous) axis, so both separable passes accumulate `C`-wide running
383    /// sums, keeping the inner channel loops contiguous. Same reassociation
384    /// caveat as the NCHW path: not bit-identical to the generic kernel.
385    unsafe fn fast_2d_separable_nhwc<T: Copy + Datum + num_traits::Float>(
386        &self,
387        input_ptr: *const T,
388        values_ptr: *mut T,
389        geo: &ConcretePoolGeometry,
390    ) where
391        usize: AsPrimitive<T>,
392    {
393        let ish = &geo.input_shape;
394        let osh = &geo.output_shape;
395        let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
396        let (ho, wo) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
397        let (kh, kw) =
398            (geo.patch.spec.kernel_shape[0] as isize, geo.patch.spec.kernel_shape[1] as isize);
399        let (pt, pl) = (geo.patch.pad_before[0] as isize, geo.patch.pad_before[1] as isize);
400        let ih_stride = *ish.h_stride() as isize;
401        let iw_stride = *ish.w_stride() as isize;
402        let oh_stride = *osh.h_stride() as isize;
403        let ow_stride = *osh.w_stride() as isize;
404        let n = *ish.n().unwrap_or(&1);
405        let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
406        let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
407        let c = *ish.c();
408
409        let axis_valid = |out: usize, k: isize, pad: isize, lim: isize| -> Vec<usize> {
410            (0..out)
411                .map(|o| {
412                    let lo = o as isize - pad;
413                    let start = (-lo).max(0);
414                    let end = (lim - lo).min(k);
415                    (end - start).max(0) as usize
416                })
417                .collect()
418        };
419        let kx_valid = axis_valid(wo, kw, pl, w);
420        let ky_valid = axis_valid(ho, kh, pt, h);
421        let full_recip: T = ((kh * kw) as usize).as_().recip();
422
423        let mut htmp = vec![T::zero(); h as usize * wo * c];
424        let mut acc = vec![T::zero(); c];
425        unsafe {
426            for nn in 0..n as isize {
427                let in_base = nn * in_stride;
428                let out_base = nn * on_stride;
429                for y in 0..h {
430                    let row = in_base + y * ih_stride;
431                    let hrow = y as usize * wo * c;
432                    acc.iter_mut().for_each(|a| *a = T::zero());
433                    for kx in 0..kw {
434                        let ix = -pl + kx;
435                        if ix >= 0 && ix < w {
436                            let p = row + ix * iw_stride;
437                            for (ch, a) in acc.iter_mut().enumerate() {
438                                *a = *a + *input_ptr.offset(p + ch as isize);
439                            }
440                        }
441                    }
442                    htmp[hrow..hrow + c].copy_from_slice(&acc);
443                    for ox in 1..wo as isize {
444                        let entering = ox - pl + kw - 1;
445                        let leaving = ox - pl - 1;
446                        if entering >= 0 && entering < w {
447                            let p = row + entering * iw_stride;
448                            for (ch, a) in acc.iter_mut().enumerate() {
449                                *a = *a + *input_ptr.offset(p + ch as isize);
450                            }
451                        }
452                        if leaving >= 0 && leaving < w {
453                            let p = row + leaving * iw_stride;
454                            for (ch, a) in acc.iter_mut().enumerate() {
455                                *a = *a - *input_ptr.offset(p + ch as isize);
456                            }
457                        }
458                        let dst = hrow + ox as usize * c;
459                        htmp[dst..dst + c].copy_from_slice(&acc);
460                    }
461                }
462                #[allow(clippy::needless_range_loop)]
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    // NNEF's `box` fragment defaults normalize to false, so this path is
539    // reachable from a model file: the sum has to actually be written.
540    #[test]
541    fn sum_pool_without_normalize_writes_the_sum() {
542        let mut model = TypedModel::default();
543        let source = model.add_source("input", f32::fact([1, 1, 4, 4])).unwrap();
544        let pool_spec = PoolSpec::new(
545            DataFormat::NCHW,
546            tvec![2, 2],
547            PaddingSpec::Valid,
548            None,
549            Some(tvec![2, 2]),
550            1,
551            1,
552        );
553        let op = SumPool { pool_spec, count_include_pad: false, normalize: false };
554        let out = model.wire_node("pool", op, &[source]).unwrap();
555        model.select_output_outlets(&out).unwrap();
556        let input = ndarray::Array4::from_shape_fn((1, 1, 4, 4), |(_, _, y, x)| (y * 4 + x) as f32)
557            .into_tensor()
558            .into_tvalue();
559        let out = model.into_runnable().unwrap().run(tvec!(input)).unwrap();
560        // windows are [0,1,4,5], [2,3,6,7], [8,9,12,13], [10,11,14,15]
561        let expected = tensor4(&[[[[10.0f32, 18.0], [42.0, 50.0]]]]);
562        out[0].close_enough(&expected, Approximation::Exact).unwrap();
563    }
564
565    #[test]
566    fn optimized_sumpool_has_concrete_geometry() {
567        let (model, input) = test_case();
568        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
569
570        let optimized = model.into_optimized().unwrap();
571        let pool = optimized
572            .nodes
573            .iter()
574            .find_map(|n| n.op_as::<OptSumPool>())
575            .expect("optimized model should contain an OptSumPool");
576        assert!(
577            pool.geometry.is_concrete(),
578            "OptSumPool geometry should be concrete after optimization"
579        );
580
581        let opt = optimized.into_runnable().unwrap().run(input).unwrap();
582        assert_eq!(*opt[0], *plain[0]);
583    }
584
585    #[test]
586    fn separable_matches_generic_kernel() {
587        let (c, h, w) = (5usize, 7usize, 9usize);
588        let pool_spec = PoolSpec::new(
589            DataFormat::NCHW,
590            tvec![3, 3],
591            PaddingSpec::SameUpper,
592            None,
593            Some(tvec![1, 1]),
594            c,
595            c,
596        );
597        let op = OptSumPool {
598            pool_spec: pool_spec.clone(),
599            count_include_pad: false,
600            normalize: true,
601            geometry: pool_spec
602                .compute_geo(&[1.to_dim(), c.to_dim(), h.to_dim(), w.to_dim()])
603                .unwrap(),
604        };
605        let input: Tensor = ndarray::Array4::from_shape_fn((1, c, h, w), |(_, cc, y, x)| {
606            ((cc * 17 + y * 3 + x) % 13) as f32 - 6.0
607        })
608        .into_tensor();
609
610        // generic zoned kernel (knob off by default)
611        let generic = op.eval(tvec![input.clone().into_tvalue()]).unwrap();
612        let generic = generic[0].try_as_plain().unwrap().as_slice::<f32>().unwrap().to_vec();
613
614        // separable kernel, called directly
615        let geo = op.geometry.to_concrete(input.shape()).unwrap();
616        let mut out = Tensor::zero::<f32>(&geo.output_shape.shape).unwrap();
617        unsafe {
618            op.fast_2d_separable::<f32>(
619                input.as_ptr::<f32>().unwrap(),
620                out.as_ptr_mut::<f32>().unwrap(),
621                geo.as_ref(),
622            );
623        }
624        let sep = out.try_as_plain().unwrap().as_slice::<f32>().unwrap();
625
626        let max_abs = generic.iter().zip(sep).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
627        assert!(max_abs < 1e-4, "separable vs generic max abs diff {max_abs}");
628    }
629
630    #[test]
631    fn separable_nhwc_matches_generic_kernel() {
632        let (c, h, w) = (5usize, 7usize, 9usize);
633        let pool_spec = PoolSpec::new(
634            DataFormat::NHWC,
635            tvec![3, 3],
636            PaddingSpec::SameUpper,
637            None,
638            Some(tvec![1, 1]),
639            c,
640            c,
641        );
642        let op = OptSumPool {
643            pool_spec: pool_spec.clone(),
644            count_include_pad: false,
645            normalize: true,
646            geometry: pool_spec
647                .compute_geo(&[1.to_dim(), h.to_dim(), w.to_dim(), c.to_dim()])
648                .unwrap(),
649        };
650        let input: Tensor = ndarray::Array4::from_shape_fn((1, h, w, c), |(_, y, x, cc)| {
651            ((cc * 17 + y * 3 + x) % 13) as f32 - 6.0
652        })
653        .into_tensor();
654
655        // generic zoned kernel (knob off by default)
656        let generic = op.eval(tvec![input.clone().into_tvalue()]).unwrap();
657        let generic = generic[0].try_as_plain().unwrap().as_slice::<f32>().unwrap().to_vec();
658
659        // separable NHWC kernel, called directly
660        let geo = op.geometry.to_concrete(input.shape()).unwrap();
661        let mut out = Tensor::zero::<f32>(&geo.output_shape.shape).unwrap();
662        unsafe {
663            op.fast_2d_separable_nhwc::<f32>(
664                input.as_ptr::<f32>().unwrap(),
665                out.as_ptr_mut::<f32>().unwrap(),
666                geo.as_ref(),
667            );
668        }
669        let sep = out.try_as_plain().unwrap().as_slice::<f32>().unwrap();
670
671        let max_abs = generic.iter().zip(sep).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
672        assert!(max_abs < 1e-4, "separable NHWC vs generic max abs diff {max_abs}");
673    }
674}
tract_core/ops/cnn/sumpool.rs

tract_core/ops/cnn/
sumpool.rs
