tract_core/ops/cnn/conv/
direct_spatial.rs

1use crate::internal::*;
2use crate::ops::cnn::Patch;
3use crate::ops::cnn::patches::{Zone, ZoneScanner};
4use crate::ops::nn::DataShape;
5
6/// A group-1 NCHW f32 convolution computed as a filter along W, four output channels at a time,
7/// instead of im2col and a matrix product. `Conv` builds it at codegen on hosts that have the
8/// `conv_w4` kernel, and the `depthwise_w` one too when the output channels are not a multiple
9/// of four.
10#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
11pub struct DirectSpatialConv {
12    patch: Patch,
13    input_shape: DataShape,
14    output_shape: DataShape,
15}
16
17impl Op for DirectSpatialConv {
18    fn name(&self) -> StaticName {
19        "DirectSpatialConv".into()
20    }
21
22    fn info(&self) -> TractResult<Vec<String>> {
23        Ok(vec![format!("{:?}", self.patch)])
24    }
25
26    fn validation(&self) -> Validation {
27        Validation::Rounding
28    }
29
30    op_as_typed_op!();
31}
32
33impl EvalOp for DirectSpatialConv {
34    op_out_of_plan!();
35
36    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
37        let (img, kernel, bias) = args_3!(inputs);
38        ensure!(img.datum_type() == f32::datum_type(), "DirectSpatialConv is f32 only");
39        let output = unsafe { self.eval_f32(&img, &kernel, &bias)? };
40        Ok(tvec!(output.into_tvalue()))
41    }
42}
43
44/// One run of output points: the kernel for the group, or a point at a time when the run is
45/// strided. Shared by the serial and the split loop structures.
46///
47/// # Safety
48/// `taps` holds `group * offsets.len()` taps in the order of `offsets`, and `ip`/`op` are
49/// valid for `len` points of the run at the strides given.
50#[allow(clippy::too_many_arguments)]
51unsafe fn one_run(
52    conv_w4: tract_linalg::routines::ConvW4F32,
53    conv_w1: Option<tract_linalg::routines::DepthwiseWF32>,
54    ip: *const f32,
55    op: *mut f32,
56    taps: &[f32],
57    offsets: &[isize],
58    b: &[f32],
59    group: usize,
60    len: usize,
61    in_stride: isize,
62    out_stride: isize,
63    oc_stride: isize,
64) {
65    unsafe {
66        if out_stride == 1 && in_stride >= 1 {
67            if group == 4 {
68                let b4 = [b[0], b[1], b[2], b[3]];
69                conv_w4(ip, op, taps, offsets, &b4, len, in_stride, oc_stride);
70            } else if let Some(conv_w1) = conv_w1 {
71                conv_w1(ip, op, taps, offsets, b[0], len, in_stride);
72            }
73            return;
74        }
75        let n_taps = offsets.len();
76        for g in 0..group {
77            let taps = &taps[g * n_taps..][..n_taps];
78            for i in 0..len as isize {
79                let ip = ip.offset(i * in_stride);
80                let sum =
81                    taps.iter().zip(offsets).fold(b[g], |sum, (t, off)| sum + t * *ip.offset(*off));
82                *op.offset(g as isize * oc_stride + i * out_stride) = sum;
83            }
84        }
85    }
86}
87
88impl DirectSpatialConv {
89    /// Every run of output points along W in `zone`: the offsets of its first input center and
90    /// first output point, its length, and the input and output steps between two points.
91    unsafe fn runs(&self, zone: &Zone, mut f: impl FnMut(isize, isize, usize, isize, isize)) {
92        if zone.valid && zone.output_ranges.len() == 2 && *self.output_shape.w_stride() == 1 {
93            let in_h = (*self.input_shape.h_stride() * self.patch.spec.strides[0]) as isize;
94            let in_w = (*self.input_shape.w_stride() * self.patch.spec.strides[1]) as isize;
95            let out_h = *self.output_shape.h_stride() as isize;
96            let x0 = zone.output_ranges[1].start as isize;
97            let len = zone.output_ranges[1].len();
98            for y in zone.output_ranges[0].clone() {
99                f(y as isize * in_h + x0 * in_w, y as isize * out_h + x0, len, in_w, 1);
100            }
101        } else {
102            let mut visitor = ZoneScanner::new(zone, &self.patch);
103            while !visitor.done {
104                f(
105                    visitor.input_center_offset,
106                    visitor.output_offset,
107                    visitor.inner_loop_len,
108                    visitor.inner_loop_input_full_stride,
109                    visitor.inner_loop_output_stride,
110                );
111                unsafe { visitor.next_non_inner_axis() };
112            }
113        }
114    }
115
116    unsafe fn eval_f32(&self, img: &Tensor, kernel: &Tensor, bias: &Tensor) -> TractResult<Tensor> {
117        let conv_w4 = tract_linalg::routines::conv_w4_f32()
118            .context("DirectSpatialConv needs a conv_w4 kernel")?;
119        let ic = *self.input_shape.c();
120        let oc = *self.output_shape.c();
121        let conv_w1 = tract_linalg::routines::depthwise_w_f32();
122        ensure!(
123            oc.is_multiple_of(4) || conv_w1.is_some(),
124            "DirectSpatialConv needs a depthwise_w kernel for its last channels"
125        );
126        let kvol = self.patch.spec.kernel_shape.iter().product::<usize>();
127        let n = *self.input_shape.n().unwrap_or(&1);
128        let n_stride_i = *self.input_shape.n_stride().unwrap_or(&0) as isize;
129        let n_stride_o = *self.output_shape.n_stride().unwrap_or(&0) as isize;
130        let c_stride_i = *self.input_shape.c_stride() as isize;
131        let c_stride_o = *self.output_shape.c_stride() as isize;
132        let mut output = unsafe { Tensor::uninitialized::<f32>(&self.output_shape.shape)? };
133        // The offsets a zone's taps read through depend on the zone and the input channels only,
134        // so they are built once per zone and shared by every output-channel group.
135        let zone_offsets: Vec<Vec<isize>> = self
136            .patch
137            .zones
138            .iter()
139            .map(|zone| {
140                let mut offsets = Vec::with_capacity(ic * kvol);
141                for c in 0..ic as isize {
142                    offsets.extend(zone.values_offsets.iter().map(|(_, off)| c * c_stride_i + off));
143                }
144                offsets
145            })
146            .collect();
147        unsafe {
148            let iptr = img.as_ptr::<f32>()?;
149            let kptr = kernel.as_ptr::<f32>()?;
150            let bptr = bias.as_ptr::<f32>()?;
151            let total = output.len();
152            let planes = n * oc;
153            let plane = total.checked_div(planes).unwrap_or(0);
154            let groups = oc / 4;
155            // A group of four output channels is four whole planes, contiguous in the output
156            // when the layout is NCHW, so a (batch, group) unit owns a block of the output that
157            // no other unit touches: the units split across the executor with no other change.
158            // Every output point keeps the arithmetic, and the value, it has serially.
159            if oc.is_multiple_of(4)
160                && groups > 0
161                && plane > 0
162                && plane * planes == total
163                && *self.output_shape.w_stride() == 1
164                && c_stride_o == plane as isize
165                && n_stride_o == (oc * plane) as isize
166            {
167                let (i_base, k_base, b_base) = (iptr as usize, kptr as usize, bptr as usize);
168                let zones = &self.patch.zones;
169                let offsets_per_zone = &zone_offsets;
170                tract_linalg::multithread::par_chunks_mut(
171                    output.as_slice_mut_unchecked::<f32>(),
172                    4 * plane,
173                    total,
174                    |first, chunk| {
175                        let mut taps = Vec::with_capacity(4 * ic * kvol);
176                        for (ix, unit_out) in chunk.chunks_mut(4 * plane).enumerate() {
177                            let unit = first + ix;
178                            let ni = (unit / groups) as isize;
179                            let o = (unit % groups) * 4;
180                            let iptr = (i_base as *const f32).offset(n_stride_i * ni);
181                            let optr = unit_out.as_mut_ptr();
182                            let b4 = [
183                                *(b_base as *const f32).add(o),
184                                *(b_base as *const f32).add(o + 1),
185                                *(b_base as *const f32).add(o + 2),
186                                *(b_base as *const f32).add(o + 3),
187                            ];
188                            for (zone_ix, zone) in zones.iter().enumerate() {
189                                let offsets = &offsets_per_zone[zone_ix];
190                                taps.clear();
191                                for g in o..o + 4 {
192                                    for c in 0..ic {
193                                        let k = (k_base as *const f32).add((g * ic + c) * kvol);
194                                        taps.extend(
195                                            zone.values_offsets.iter().map(|(ix, _)| *k.add(*ix)),
196                                        );
197                                    }
198                                }
199                                self.runs(zone, |center, out, len, in_stride, out_stride| {
200                                    one_run(
201                                        conv_w4,
202                                        conv_w1,
203                                        iptr.offset(center),
204                                        optr.offset(out),
205                                        &taps,
206                                        offsets,
207                                        &b4,
208                                        4,
209                                        len,
210                                        in_stride,
211                                        out_stride,
212                                        c_stride_o,
213                                    )
214                                });
215                            }
216                        }
217                        Ok(())
218                    },
219                )?;
220                return Ok(output);
221            }
222            let optr = output.as_ptr_mut::<f32>()?;
223            let mut taps = Vec::with_capacity(4 * ic * kvol);
224            for (zone_ix, zone) in self.patch.zones.iter().enumerate() {
225                let offsets = &zone_offsets[zone_ix];
226                let mut o = 0;
227                while o < oc {
228                    let group = if o + 4 <= oc { 4 } else { 1 };
229                    taps.clear();
230                    for g in o..o + group {
231                        for c in 0..ic {
232                            let k = kptr.add((g * ic + c) * kvol);
233                            taps.extend(zone.values_offsets.iter().map(|(ix, _)| *k.add(*ix)));
234                        }
235                    }
236                    let b = std::slice::from_raw_parts(bptr.add(o), group);
237                    for ni in 0..n as isize {
238                        let iptr = iptr.offset(n_stride_i * ni);
239                        let optr = optr.offset(n_stride_o * ni + c_stride_o * o as isize);
240                        self.runs(zone, |center, out, len, in_stride, out_stride| {
241                            one_run(
242                                conv_w4,
243                                conv_w1,
244                                iptr.offset(center),
245                                optr.offset(out),
246                                &taps,
247                                offsets,
248                                b,
249                                group,
250                                len,
251                                in_stride,
252                                out_stride,
253                                c_stride_o,
254                            )
255                        });
256                    }
257                    o += group;
258                }
259            }
260        }
261        Ok(output)
262    }
263}
264
265impl TypedOp for DirectSpatialConv {
266    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
267        ensure!(inputs.len() == 3);
268        Ok(tvec!(inputs[0].datum_type.fact(&self.output_shape.shape)))
269    }
270
271    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
272        let [_input, kernel, _bias] = inputs else {
273            bail!("DirectSpatialConv expects three inputs");
274        };
275        let n_output_points = self.patch.output_shape.iter().cloned().product::<usize>();
276        Ok(tvec!((
277            Cost::FMA(inputs[0].datum_type),
278            kernel.shape.volume() * self.input_shape.n().unwrap_or(&1) * n_output_points
279        )))
280    }
281
282    as_op!();
283}
tract_core/ops/cnn/conv/direct_spatial.rs

tract_core/ops/cnn/conv/
direct_spatial.rs
