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

tract_core/ops/cnn/
sumpool.rs
