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

tract_core/ops/cnn/
sumpool.rs
