tract_core/ops/cnn/conv/
mod.rs

1mod block_quant;
2mod blocked;
3#[allow(clippy::module_inception)]
4mod conv;
5mod depth_wise;
6mod direct_spatial;
7mod im2col;
8mod lazy_im2col;
9mod nhwc;
10mod q_sum_b;
11
12use crate::internal::*;
13use crate::ops::cnn::Deconv;
14
15pub use self::blocked::BlockedConv;
16pub use self::conv::Conv;
17pub use self::direct_spatial::DirectSpatialConv;
18pub use self::im2col::Im2Col;
19pub(crate) use self::q_sum_b::QSumB;
20
21#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)]
22pub enum KernelFormat {
23    #[default]
24    OIHW,
25    HWIO,
26    OHWI,
27}
28
29impl KernelFormat {
30    pub fn h_axis(&self) -> usize {
31        match self {
32            KernelFormat::OIHW => 2,
33            KernelFormat::HWIO => 0,
34            KernelFormat::OHWI => 1,
35        }
36    }
37
38    pub fn spatial_shape<'a, D>(&self, full_shape: &'a [D]) -> &'a [D] {
39        &full_shape[self.h_axis()..][..full_shape.len() - 2]
40    }
41
42    pub fn hw<'a, D>(&self, full_shape: &'a [D]) -> &'a [D] {
43        self.spatial_shape(full_shape)
44    }
45
46    pub fn i<'a, D>(&self, full_shape: &'a [D]) -> &'a D {
47        match self {
48            KernelFormat::OIHW => &full_shape[1],
49            KernelFormat::HWIO => &full_shape[full_shape.len() - 2],
50            KernelFormat::OHWI => &full_shape[full_shape.len() - 1],
51        }
52    }
53
54    pub fn o_axis<D>(&self, full_shape: &[D]) -> usize {
55        match self {
56            KernelFormat::OIHW | KernelFormat::OHWI => 0,
57            KernelFormat::HWIO => full_shape.len() - 1,
58        }
59    }
60
61    pub fn i_axis<D>(&self, full_shape: &[D]) -> usize {
62        match self {
63            KernelFormat::OIHW => 1,
64            KernelFormat::OHWI => full_shape.len() - 1,
65            KernelFormat::HWIO => full_shape.len() - 2,
66        }
67    }
68
69    pub fn o<'a, D>(&self, full_shape: &'a [D]) -> &'a D {
70        &full_shape[self.o_axis(full_shape)]
71    }
72
73    pub fn input_channels<'s, D: DimLike>(
74        &self,
75        full_kernel_shape: &'s [D],
76        group: usize,
77    ) -> Cow<'s, D> {
78        match self {
79            KernelFormat::OIHW => Cow::Owned(self.i(full_kernel_shape).clone() * group),
80            KernelFormat::HWIO | KernelFormat::OHWI => Cow::Borrowed(self.i(full_kernel_shape)),
81        }
82    }
83
84    pub fn output_channels<'s, D: DimLike>(
85        &self,
86        full_kernel_shape: &'s [D],
87        group: usize,
88    ) -> Cow<'s, D> {
89        match self {
90            KernelFormat::OIHW => Cow::Borrowed(self.o(full_kernel_shape)),
91            KernelFormat::HWIO | KernelFormat::OHWI => {
92                Cow::Owned(self.o(full_kernel_shape).clone() * group)
93            }
94        }
95    }
96
97    pub fn kernel_as_group_o_i_h_w_ops(
98        &self,
99        full_shape: &[impl DimLike],
100        group: usize,
101    ) -> TVec<AxisOp> {
102        let geo_rank = full_shape.len() - 2;
103        match self {
104            // g is on i
105            KernelFormat::HWIO => {
106                tvec!(
107                    AxisOp::Reshape(
108                        geo_rank,
109                        tvec!(self.i(full_shape).to_dim()),
110                        tvec!(group.to_dim(), self.i(full_shape).to_dim() / group),
111                    ), // h w g i o
112                    AxisOp::Move(geo_rank, 0),     // g h w i o
113                    AxisOp::Move(geo_rank + 2, 1), // g o h w i
114                    AxisOp::Move(geo_rank + 2, 2)
115                ) // g o i h w
116            }
117            // g is on o
118            KernelFormat::OIHW => {
119                tvec!(AxisOp::Reshape(
120                    0,
121                    tvec!(self.o(full_shape).to_dim()),
122                    tvec!(group.to_dim(), self.o(full_shape).to_dim() / group),
123                ))
124            }
125            // g is on i
126            KernelFormat::OHWI => {
127                tvec!(
128                    AxisOp::Reshape(
129                        geo_rank + 1,
130                        tvec!(self.i(full_shape).to_dim()),
131                        tvec!(group.to_dim(), self.i(full_shape).to_dim() / group),
132                    ), // o h w g i
133                    AxisOp::Move(geo_rank + 1, 0), // g o h w i
134                    AxisOp::Move(geo_rank + 2, 2)
135                )
136            }
137        }
138    }
139
140    pub fn kernel_as_group_o_i_hw_ops(
141        &self,
142        full_shape: &[impl DimLike],
143        group: usize,
144    ) -> TVec<AxisOp> {
145        let mut ops = self.kernel_as_group_o_i_h_w_ops(full_shape, group);
146        if self.hw(full_shape).len() > 1 {
147            ops.push(AxisOp::Reshape(
148                3,
149                self.hw(full_shape).iter().map(|t| t.to_dim()).collect(),
150                tvec!(self.hw(full_shape).iter().map(|t| t.to_dim()).product()),
151            ));
152        }
153        ops
154    }
155
156    pub fn kernel_as_group_o_ihw_ops(
157        &self,
158        full_shape: &[impl DimLike],
159        group: usize,
160    ) -> TVec<AxisOp> {
161        let i = (self.input_channels(full_shape, group).into_owned() / group).to_dim();
162        let hw = self.hw(full_shape).iter().map(|t| t.to_dim()).product::<TDim>();
163        let mut ops = self.kernel_as_group_o_i_hw_ops(full_shape, group);
164        ops.push(AxisOp::Reshape(2, tvec!(i.clone(), hw.clone()), tvec!(i * hw)));
165        ops
166    }
167
168    pub fn kernel_as_group_o_i_hw(&self, kernel: &Tensor, group: usize) -> TractResult<Tensor> {
169        let mut kernel = kernel.clone();
170        let ops = self.kernel_as_group_o_i_hw_ops(kernel.shape(), group);
171        for op in &ops {
172            op.change_tensor(&mut kernel, false)?;
173        }
174        Ok(kernel)
175    }
176
177    pub fn kernel_as_group_o_ihw(&self, kernel: &Tensor, group: usize) -> TractResult<Tensor> {
178        let group_o_i_hw = self.kernel_as_group_o_i_hw(kernel, group)?;
179        Ok(group_o_i_hw.collapse_axis_with_next(2))
180    }
181}
182
183pub fn rewrite_kernel_conv_in_oihw(
184    _ctx: &(),
185    model: &TypedModel,
186    node: &TypedNode,
187    name: &str,
188    conv: &Conv,
189) -> TractResult<Option<TypedModelPatch>> {
190    rewrite_kernel_in_oihw(
191        model,
192        node,
193        name,
194        conv.kernel_fmt,
195        conv.group,
196        Box::new(Conv { kernel_fmt: KernelFormat::OIHW, ..conv.clone() }),
197    )
198}
199
200pub fn rewrite_kernel_deconv_in_oihw(
201    _ctx: &(),
202    model: &TypedModel,
203    node: &TypedNode,
204    name: &str,
205    conv: &Deconv,
206) -> TractResult<Option<TypedModelPatch>> {
207    rewrite_kernel_in_oihw(
208        model,
209        node,
210        name,
211        conv.kernel_format,
212        conv.group,
213        Box::new(Deconv { kernel_format: KernelFormat::OIHW, ..conv.clone() }),
214    )
215}
216
217fn rewrite_kernel_in_oihw(
218    model: &TypedModel,
219    node: &TypedNode,
220    name: &str,
221    fmt: KernelFormat,
222    group: usize,
223    new: Box<dyn TypedOp>,
224) -> TractResult<Option<TypedModelPatch>> {
225    rule_if!(fmt != KernelFormat::OIHW);
226    let mut patch = TypedModelPatch::default();
227    let mut wire = patch.taps(model, &node.inputs)?;
228    let prefix = format!("{name}.kernel_reorg");
229    for (ix, op) in fmt
230        .kernel_as_group_o_i_h_w_ops(&patch.outlet_fact(wire[1])?.shape, group)
231        .into_iter()
232        .enumerate()
233    {
234        wire[1] = patch.wire_node(format!("{prefix}.{ix}"), op, &[wire[1]])?[0];
235    }
236    wire[1] =
237        AxisOp::wire_collapse_axis(&mut patch, format!("{name}.kernel_reorg_go"), wire[1], 0)?[0];
238    wire = patch.wire_node(name, new, &wire)?;
239    patch.shunt_outside(model, node.id.into(), wire[0])?;
240    Ok(Some(patch))
241}
tract_core/ops/cnn/conv/mod.rs

tract_core/ops/cnn/conv/
mod.rs
