#![allow(dead_code)]
use std::ops::AddAssign;
use crate::internal::*;
use crate::ops::cnn::padding::ComputedPaddedDim;
use crate::ops::cnn::{KernelFormat, PoolSpec};
use crate::ops::nn::DataShape;
use tract_ndarray::prelude::*;
use tract_num_traits::Float;
#[derive(Clone, Debug, new, Hash, PartialEq, Eq)]
pub struct DeconvSum {
pub pool_spec: PoolSpec,
pub kernel_format: KernelFormat,
pub input_shape: ShapeFact,
pub adjustments: TVec<usize>,
pub group: usize,
}
impl Op for DeconvSum {
fn name(&self) -> StaticName {
"DeconvSum".into()
}
op_as_typed_op!();
}
impl EvalOp for DeconvSum {
op_out_of_plan!();
fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
self.eval_with_values(inputs, ctx.symbols)
}
}
impl DeconvSum {
fn eval_with_values(
&self,
inputs: TVec<TValue>,
values: &SymbolValues,
) -> TractResult<TVec<TValue>> {
let (gemm, bias) = args_2!(inputs);
let input_shape = self.input_shape.eval_to_usize(values)?.into_owned();
let input_shape = self.pool_spec.data_format.shape(input_shape)?;
let output_shape =
super::output_shape(&self.pool_spec, &input_shape.shape, &self.adjustments)?;
let output_shape = self.pool_spec.data_format.shape(output_shape)?;
let spatial_output_details = self.pool_spec.padding.compute_for_deconv(
input_shape.hw_dims(),
&self.pool_spec.kernel_shape,
&self.pool_spec.dilations(),
&self.pool_spec.strides(),
&self.adjustments,
)?;
let hw = *gemm.shape().last().unwrap();
let n = *output_shape.n().unwrap_or(&1);
let n_o_hkwk_hw = gemm.into_tensor().into_shape(&[
n,
*output_shape.c(),
self.pool_spec.kernel_shape.iter().product(),
hw,
])?;
let per_channel_bias = bias.len() == 1 || bias.len() == *output_shape.c();
let geometry =
non_overlapping_geometry(self, &input_shape, &output_shape, &spatial_output_details);
let mut tensor = if let (Some((input_hw, strides)), true) = (geometry, per_channel_bias) {
let mut tensor = Tensor::zero_dt(n_o_hkwk_hw.datum_type(), &output_shape.shape)?;
if !self.pool_spec.data_format.has_n() {
tensor.insert_axis(0)?;
}
eval_non_overlapping(
&n_o_hkwk_hw,
&bias,
&mut tensor,
*output_shape.c(),
input_hw,
strides,
self.pool_spec.data_format.c_is_last(),
)?;
tensor
} else {
let mut tensor = if bias.shape() == &*output_shape.shape {
bias.into_tensor()
} else {
bias.broadcast_to_shape(&output_shape.shape)?
};
if !self.pool_spec.data_format.has_n() {
tensor.insert_axis(0)?;
}
eval(
self,
&input_shape,
&output_shape,
&spatial_output_details,
&n_o_hkwk_hw,
&mut tensor,
)?;
tensor
};
if !self.pool_spec.data_format.has_n() {
tensor.remove_axis(0)?;
}
Ok(tvec!(tensor.into_tvalue()))
}
}
impl TypedOp for DeconvSum {
fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
ensure!(inputs.len() == 2);
let shape = super::output_shape(&self.pool_spec, &self.input_shape, &self.adjustments)?;
ensure!(inputs[1].rank() == shape.len());
ensure!(
inputs[1].shape.iter().zip(shape.iter()).all(|(b, o)| b.is_one() || b == o),
"DeconvSum bias {:?} does not broadcast to output {:?}",
inputs[1].shape,
shape
);
Ok(tvec!(inputs[0].datum_type.fact(shape)))
}
fn set_symbols(
&self,
_source: &TypedModel,
node: &TypedNode,
target: &mut TypedModel,
mapping: &HashMap<OutletId, OutletId>,
subs: &HashMap<Symbol, TDim>,
) -> TractResult<TVec<OutletId>> {
target.wire_node(
&node.name,
Self { input_shape: self.input_shape.substitute(subs)?.into_owned(), ..self.clone() },
&[mapping[&node.inputs[0]], mapping[&node.inputs[1]]],
)
}
as_op!();
}
fn non_overlapping_geometry(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
) -> Option<([usize; 2], [usize; 2])> {
let rank = input_shape.hw_rank();
if rank == 0 || rank > 2 {
return None;
}
let strides = op.pool_spec.strides();
let dilations = op.pool_spec.dilations();
let input_hw = input_shape.hw_dims();
for ax in 0..rank {
if op.pool_spec.kernel_shape[ax] != strides[ax]
|| dilations[ax] != 1
|| spatial_output_details[ax].pad_before != 0
|| output_shape.hw_dims()[ax] != input_hw[ax] * strides[ax]
{
return None;
}
}
Some(if rank == 1 {
([1, input_hw[0]], [1, strides[0]])
} else {
([input_hw[0], input_hw[1]], [strides[0], strides[1]])
})
}
fn eval_non_overlapping(
gemm: &Tensor,
bias: &Tensor,
output: &mut Tensor,
channels: usize,
input_hw: [usize; 2],
strides: [usize; 2],
c_is_last: bool,
) -> TractResult<()> {
ensure!(bias.len() == 1 || bias.len() == channels);
let dt = output.datum_type();
dispatch_floatlike!(write_non_overlapping(dt)(
gemm,
bias,
output,
channels,
input_hw,
strides,
c_is_last,
|a, b| a + b
))
}
#[allow(clippy::too_many_arguments)]
fn write_non_overlapping<T: Datum + Copy>(
gemm: &Tensor,
bias: &Tensor,
output: &mut Tensor,
channels: usize,
input_hw: [usize; 2],
strides: [usize; 2],
c_is_last: bool,
add: impl Fn(T, T) -> T + Copy,
) -> TractResult<()> {
let gemm_view = gemm.try_as_plain_ram()?;
let gemm = gemm_view.as_slice::<T>()?;
let bias_view = bias.try_as_plain_ram()?;
let bias = bias_view.as_slice::<T>()?;
let mut output_view = output.try_as_plain_ram_mut()?;
let output = output_view.as_slice_mut::<T>()?;
let bias_at = |c: usize| if bias.len() == 1 { bias[0] } else { bias[c] };
let [ih, iw] = input_hw;
let [sh, sw] = strides;
let (oh, ow) = (ih * sh, iw * sw);
let plane = ih * iw;
let taps = sh * sw;
if c_is_last {
for (gemm_n, output_n) in gemm
.chunks_exact(channels * taps * plane)
.zip(output.chunks_exact_mut(oh * ow * channels))
{
for (oy, row) in output_n.chunks_exact_mut(ow * channels).enumerate() {
let (iy, ky) = (oy / sh, oy % sh);
for (ox, pixel) in row.chunks_exact_mut(channels).enumerate() {
let (ix, kx) = (ox / sw, ox % sw);
let tap = (ky * sw + kx) * plane + iy * iw + ix;
for (c, out) in pixel.iter_mut().enumerate() {
*out = add(gemm_n[c * taps * plane + tap], bias_at(c));
}
}
}
}
} else {
for (nc, (gemm_nc, output_nc)) in
gemm.chunks_exact(taps * plane).zip(output.chunks_exact_mut(oh * ow)).enumerate()
{
let b = bias_at(nc % channels);
for (oy, row) in output_nc.chunks_exact_mut(ow).enumerate() {
let (iy, ky) = (oy / sh, oy % sh);
let tap_row = |kx: usize| &gemm_nc[(ky * sw + kx) * plane + iy * iw..][..iw];
if sw == 2 {
for (pair, (&even, &odd)) in
row.chunks_exact_mut(2).zip(tap_row(0).iter().zip(tap_row(1)))
{
pair[0] = add(even, b);
pair[1] = add(odd, b);
}
} else {
for kx in 0..sw {
for (out, &g) in row[kx..].iter_mut().step_by(sw).zip(tap_row(kx)) {
*out = add(g, b);
}
}
}
}
}
}
Ok(())
}
fn eval(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
n_o_hkwk_hw: &Tensor,
output: &mut Tensor,
) -> TractResult<()> {
let dt = output.datum_type();
unsafe {
#[cfg(target_arch = "aarch64")]
if dt == f16::datum_type() && tract_linalg::arm64::has_fp16() {
return eval_t_aarch64fp16::<f16>(
op,
input_shape,
output_shape,
spatial_output_details,
n_o_hkwk_hw,
output,
|a, b| tract_linalg::arm64::add_f16(a, b),
);
}
dispatch_floatlike!(eval_t_generic(dt)(
op,
input_shape,
output_shape,
spatial_output_details,
n_o_hkwk_hw,
output,
|a, b| a + b
))
}
}
macro_rules! impl_eval {
($(#[$meta: meta])* $suffix: ident) => {
pastey::paste! {
$(#[$meta])*
unsafe fn [<eval_t_ $suffix>]<T: Datum + Float + Copy + AddAssign<T>>(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
n_o_hkwk_hw: &Tensor,
output: &mut Tensor,
add: impl Fn(T, T) -> T + Copy + 'static,
) -> TractResult<()> {
let mut output_plain = output.try_as_plain_ram_mut()?;
let output = output_plain.to_array_view_mut::<T>()?;
let n_o_hkwk_hw: ArrayView4<T> = n_o_hkwk_hw.to_plain_array_view::<T>()?.into_dimensionality()?;
match input_shape.hw_rank() {
1 => [<main_loop_1d_ $suffix>](
op,
input_shape,
output_shape,
spatial_output_details,
&n_o_hkwk_hw,
&mut output.into_dimensionality().unwrap(),
add,
),
2 => [<main_loop_2d_ $suffix>](
op,
input_shape,
output_shape,
spatial_output_details,
&n_o_hkwk_hw,
&mut output.into_dimensionality().unwrap(),
add,
),
3 => [<main_loop_3d_ $suffix>](
op,
input_shape,
output_shape,
spatial_output_details,
&n_o_hkwk_hw,
&mut output.into_dimensionality().unwrap(),
add,
),
_ => [<main_loop_ $suffix>](
op,
input_shape,
output_shape,
spatial_output_details,
&n_o_hkwk_hw,
&mut output.into_dimensionality().unwrap(),
add,
),
}
}
pub fn [<main_loop_1d_ $suffix>]<T: Datum + Float>(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
n_o_hkwk_hw: &ArrayView4<T>,
output: &mut ArrayViewMut3<T>,
add: impl Fn(T, T) -> T + Copy + 'static,
) -> TractResult<()> {
let n = *output_shape.n().unwrap_or(&1);
let kernel_len = op.pool_spec.kernel_shape[0];
let geo_input_len = input_shape.hw_dims()[0];
let geo_output_len = output_shape.hw_dims()[0];
let x_stride = op.pool_spec.strides().as_ref()[0];
let x_dil = op.pool_spec.dilations().as_ref()[0];
let x_pad = spatial_output_details[0].pad_before as isize;
for n in 0..n {
for o in 0..*output_shape.c() {
for kx in 0..kernel_len {
for gx in 0..geo_input_len {
let x = (kx * x_dil + gx * x_stride) as isize - x_pad;
if x < 0 || x >= geo_output_len as isize {
continue;
}
let coord = if op.pool_spec.data_format.c_is_last() {
[n, x as usize, o]
} else {
[n, o, x as usize]
};
unsafe {
let value = *n_o_hkwk_hw.uget((n, o, kx, gx));
*output.uget_mut(coord) = add(*output.uget(coord), value);
}
}
}
}
}
Ok(())
}
pub fn [<main_loop_2d_ $suffix>]<T: Datum + Float>(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
n_o_hkwk_hw: &ArrayView4<T>,
output: &mut ArrayViewMut4<T>,
add: impl Fn(T, T) -> T + Copy + 'static,
) -> TractResult<()> {
let n = *output_shape.n().unwrap_or(&1);
let x_stride = op.pool_spec.strides().as_ref()[0];
let y_stride = op.pool_spec.strides().as_ref()[1];
let x_dil = op.pool_spec.dilations().as_ref()[0];
let y_dil = op.pool_spec.dilations().as_ref()[1];
let x_pad = spatial_output_details[0].pad_before as isize;
let y_pad = spatial_output_details[1].pad_before as isize;
let output_c = *output_shape.c();
let output_c_stride = *output_shape.c_stride() as isize;
let output_x_stride = output_shape.hw_strides()[0] as isize;
let output_y_stride = output_shape.hw_strides()[1] as isize;
let temp_n_stride = n_o_hkwk_hw.strides()[0];
let temp_o_stride = n_o_hkwk_hw.strides()[1];
let temp_k_stride = n_o_hkwk_hw.strides()[2];
let temp_i_stride = n_o_hkwk_hw.strides()[3];
let ox_len = output_shape.hw_dims()[0];
let oy_len = output_shape.hw_dims()[1];
let ix_len = input_shape.hw_dims()[0];
let iy_len = input_shape.hw_dims()[1];
let kx_len = op.pool_spec.kernel_shape[0];
let ky_len = op.pool_spec.kernel_shape[1];
if !op.pool_spec.data_format.c_is_last()
&& iy_len >= 16
&& let Some(temp) = n_o_hkwk_hw.as_slice()
&& let Some(output) = output.as_slice_mut()
{
for n in 0..n {
for o in 0..output_c {
let output_base = n * *output_shape.n_stride().unwrap_or(&0)
+ o * output_c_stride as usize;
let temp_base = n * temp_n_stride as usize + o * temp_o_stride as usize;
for kx in 0..kx_len {
let x_base = (kx * x_dil) as isize - x_pad;
let ix_start = ((-x_base).max(0) as usize).div_ceil(x_stride);
let ix_end = ((ox_len as isize - x_base).max(0) as usize).div_ceil(x_stride).min(ix_len);
for ky in 0..ky_len {
let y_base = (ky * y_dil) as isize - y_pad;
let iy_start = ((-y_base).max(0) as usize).div_ceil(y_stride);
let iy_end = ((oy_len as isize - y_base).max(0) as usize).div_ceil(y_stride).min(iy_len);
if iy_start >= iy_end {
continue;
}
for ix in ix_start..ix_end {
let ox = (x_base + (ix * x_stride) as isize) as usize;
let oy = (y_base + (iy_start * y_stride) as isize) as usize;
let input_start = temp_base + (kx * ky_len + ky) * temp_k_stride as usize
+ ix * iy_len + iy_start;
let output_start = output_base + ox * output_x_stride as usize + oy;
let input = &temp[input_start..input_start + iy_end - iy_start];
let output = &mut output[output_start..output_start + (input.len() - 1) * y_stride + 1];
if y_stride == 1 {
for (output, &input) in output.iter_mut().zip(input) {
*output = add(*output, input);
}
} else {
let blocks = (input.len() - 1) / 8;
let (input_head, input_tail) = input.split_at(blocks * 8);
let (output_head, output_tail) = output.split_at_mut(blocks * 8 * y_stride);
for (output, input) in output_head.chunks_exact_mut(8 * y_stride).zip(input_head.chunks_exact(8)) {
for i in 0..8 {
output[i * y_stride] = add(output[i * y_stride], input[i]);
}
}
for (output, &input) in output_tail.iter_mut().step_by(y_stride).zip(input_tail) {
*output = add(*output, input);
}
}
}
}
}
}
}
return Ok(());
}
unsafe {
for n in 0..n {
let output = output.as_mut_ptr().add(n * *output_shape.n_stride().unwrap_or(&0));
let temp = n_o_hkwk_hw.as_ptr().offset(n as isize * temp_n_stride);
for kx in 0..kx_len {
let temp = temp.offset((kx * ky_len) as isize * temp_k_stride);
for ix in 0..ix_len {
let ox = (kx * x_dil + ix * x_stride) as isize - x_pad;
if ox < 0 || ox >= ox_len as isize {
continue;
}
let temp = temp.offset((ix * iy_len) as isize * temp_i_stride);
let output = output.offset(ox * output_x_stride);
for ky in 0..ky_len {
let temp = temp.offset(ky as isize * temp_k_stride);
let oy = (ky * y_dil) as isize - y_pad;
for iy in 0..iy_len {
let oy = oy + (iy * y_stride) as isize;
if oy < 0 || oy >= oy_len as isize {
continue;
}
let temp = temp.offset(iy as isize * temp_i_stride);
let output = output.offset(oy * output_y_stride);
[<main_loop_2d_inner_ $suffix>](
output_c,
temp,
temp_o_stride,
output,
output_c_stride,
add,
)
}
}
}
}
}
}
Ok(())
}
#[inline(never)]
#[allow(clippy::erasing_op)]
#[allow(clippy::identity_op)]
unsafe fn [<main_loop_2d_inner_ $suffix>]<T: Datum + Float>(
output_c: usize,
temp: *const T,
temp_o_stride: isize,
output: *mut T,
output_c_stride: isize,
add: impl Fn(T, T) -> T + Copy + 'static,
) { unsafe {
let mut c = 0;
let mut right = temp;
let mut left = output;
while c + 8 < output_c {
let mut left0 = *left.offset(0 * output_c_stride);
let mut left1 = *left.offset(1 * output_c_stride);
let mut left2 = *left.offset(2 * output_c_stride);
let mut left3 = *left.offset(3 * output_c_stride);
let mut left4 = *left.offset(4 * output_c_stride);
let mut left5 = *left.offset(5 * output_c_stride);
let mut left6 = *left.offset(6 * output_c_stride);
let mut left7 = *left.offset(7 * output_c_stride);
let right0 = *right.offset(0 * temp_o_stride);
let right1 = *right.offset(1 * temp_o_stride);
let right2 = *right.offset(2 * temp_o_stride);
let right3 = *right.offset(3 * temp_o_stride);
let right4 = *right.offset(4 * temp_o_stride);
let right5 = *right.offset(5 * temp_o_stride);
let right6 = *right.offset(6 * temp_o_stride);
let right7 = *right.offset(7 * temp_o_stride);
left0 = add(left0, right0);
left1 = add(left1, right1);
left2 = add(left2, right2);
left3 = add(left3, right3);
left4 = add(left4, right4);
left5 = add(left5, right5);
left6 = add(left6, right6);
left7 = add(left7, right7);
*left.offset(0 * output_c_stride) = left0;
*left.offset(1 * output_c_stride) = left1;
*left.offset(2 * output_c_stride) = left2;
*left.offset(3 * output_c_stride) = left3;
*left.offset(4 * output_c_stride) = left4;
*left.offset(5 * output_c_stride) = left5;
*left.offset(6 * output_c_stride) = left6;
*left.offset(7 * output_c_stride) = left7;
c += 8;
left = left.offset(8 * output_c_stride);
right = right.offset(8 * temp_o_stride);
}
for c in c..output_c {
let value = *temp.offset(c as isize * temp_o_stride);
let ptr = output.offset(c as isize * output_c_stride);
*ptr = add(*ptr, value);
}
}}
pub fn [<main_loop_3d_ $suffix>]<T: Datum + Float>(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
n_o_hkwk_hw: &ArrayView4<T>,
output: &mut ArrayViewMut5<T>,
add: impl Fn(T, T) -> T + Copy + 'static,
) -> TractResult<()> {
let n = *output_shape.n().unwrap_or(&1);
let kernel_shape: [usize; 3] =
[op.pool_spec.kernel_shape[0], op.pool_spec.kernel_shape[1], op.pool_spec.kernel_shape[2]];
let geo_input_shape: [usize; 3] =
[input_shape.hw_dims()[0], input_shape.hw_dims()[1], input_shape.hw_dims()[2]];
let geo_output_shape: [usize; 3] =
[output_shape.hw_dims()[0], output_shape.hw_dims()[1], output_shape.hw_dims()[2]];
let x_stride = op.pool_spec.strides().as_ref()[0];
let y_stride = op.pool_spec.strides().as_ref()[1];
let z_stride = op.pool_spec.strides().as_ref()[2];
let x_dil = op.pool_spec.dilations().as_ref()[0];
let y_dil = op.pool_spec.dilations().as_ref()[1];
let z_dil = op.pool_spec.dilations().as_ref()[2];
let x_pad = spatial_output_details[0].pad_before as isize;
let y_pad = spatial_output_details[1].pad_before as isize;
let z_pad = spatial_output_details[2].pad_before as isize;
for n in 0..n {
for o in 0..*output_shape.c() {
for (kix, (kx, ky, kz)) in tract_ndarray::indices(kernel_shape).into_iter().enumerate()
{
for (gix, (gx, gy, gz)) in
tract_ndarray::indices(geo_input_shape).into_iter().enumerate()
{
let x = (kx * x_dil + gx * x_stride) as isize - x_pad;
let y = (ky * y_dil + gy * y_stride) as isize - y_pad;
let z = (kz * z_dil + gz * z_stride) as isize - z_pad;
if x < 0
|| y < 0
|| z < 0
|| x >= geo_output_shape[0] as isize
|| y >= geo_output_shape[1] as isize
|| z >= geo_output_shape[2] as isize
{
continue;
}
let coord = if op.pool_spec.data_format.c_is_last() {
[n, x as usize, y as usize, z as usize, o]
} else {
[n, o, x as usize, y as usize, z as usize]
};
unsafe {
let value = *n_o_hkwk_hw.uget((n, o, kix, gix));
*output.uget_mut(coord) = add(*output.uget(coord), value);
}
}
}
}
}
Ok(())
}
pub fn [<main_loop_ $suffix>]<T: Datum + Float>(
op: &DeconvSum,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
n_o_hkwk_hw: &ArrayView4<T>,
output: &mut ArrayViewMutD<T>,
add: impl Fn(T, T) -> T + Copy + 'static,
) -> TractResult<()> {
let n = *output_shape.n().unwrap_or(&1);
let strides = op.pool_spec.strides();
let dilations = op.pool_spec.dilations();
for n in 0..n {
for o in 0..*output_shape.c() {
for (kix, kcoords) in
tract_ndarray::indices(&*op.pool_spec.kernel_shape).into_iter().enumerate()
{
for (gix, gcoords) in
tract_ndarray::indices(input_shape.hw_dims()).into_iter().enumerate()
{
let ocoord: TVec<isize> = tract_itertools::izip!(
kcoords.slice(),
gcoords.slice(),
strides.as_ref(),
dilations.as_ref(),
spatial_output_details
)
.map(|(k, g, s, d, details)| {
(k * d + g * s) as isize - details.pad_before as isize
})
.collect();
if ocoord
.iter()
.zip(output_shape.hw_dims().iter())
.all(|(x, dim)| *x >= 0 && (*x as usize) < *dim)
{
let ocoord = ocoord.iter().map(|x| *x as usize).collect::<TVec<_>>();
let ocoord = op.pool_spec.data_format.with_n().from_n_c_hw(n, o, ocoord)?;
let value = n_o_hkwk_hw[(n, o, kix, gix)];
output[&*ocoord.shape] = add(output[&*ocoord.shape], value)
}
}
}
}
}
Ok(())
}
}
}
}
impl_eval!(generic);
impl_eval! {
#[target_feature(enable = "fp16")]
#[cfg(target_arch = "aarch64")]
aarch64fp16
}
crate::declare_knob!(
scopable TRACT_DISABLE_DEPTHWISE_DECONV,
bool,
false,
"Disable the fused depthwise-deconvolution path, falling back to einsum + DeconvSum."
);
#[derive(Clone, Debug, new, Hash, PartialEq, Eq)]
pub struct DepthwiseDeconv {
pub pool_spec: PoolSpec,
pub kernel_format: KernelFormat,
pub input_shape: ShapeFact,
pub adjustments: TVec<usize>,
pub group: usize,
}
impl Op for DepthwiseDeconv {
fn name(&self) -> StaticName {
"DepthwiseDeconv".into()
}
op_as_typed_op!();
}
impl EvalOp for DepthwiseDeconv {
op_out_of_plan!();
fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
let (kernel, input, bias) = args_3!(inputs);
let input_shape = self.input_shape.eval_to_usize(ctx.symbols)?.into_owned();
let input_shape = self.pool_spec.data_format.shape(input_shape)?;
let output_shape =
super::output_shape(&self.pool_spec, &input_shape.shape, &self.adjustments)?;
let output_shape = self.pool_spec.data_format.shape(output_shape)?;
let spatial_output_details = self.pool_spec.padding.compute_for_deconv(
input_shape.hw_dims(),
&self.pool_spec.kernel_shape,
&self.pool_spec.dilations(),
&self.pool_spec.strides(),
&self.adjustments,
)?;
let mut tensor = bias.into_tensor();
if !self.pool_spec.data_format.has_n() {
tensor.insert_axis(0)?;
}
dispatch_floatlike!(Self::eval_t(tensor.datum_type())(
self,
&input_shape,
&output_shape,
&spatial_output_details,
&kernel,
&input,
&mut tensor
))?;
if !self.pool_spec.data_format.has_n() {
tensor.remove_axis(0)?;
}
Ok(tvec!(tensor.into_tvalue()))
}
}
impl DepthwiseDeconv {
#[allow(clippy::too_many_arguments)]
unsafe fn eval_t_2d<T: Datum + Float + Copy + AddAssign<T>>(
&self,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
kernel: &[T],
input: &[T],
output: &mut Tensor,
) -> TractResult<()> {
let n_batch = *output_shape.n().unwrap_or(&1);
let x_stride = self.pool_spec.strides().as_ref()[0];
let y_stride = self.pool_spec.strides().as_ref()[1];
let x_dil = self.pool_spec.dilations().as_ref()[0];
let y_dil = self.pool_spec.dilations().as_ref()[1];
let x_pad = spatial_output_details[0].pad_before as isize;
let y_pad = spatial_output_details[1].pad_before as isize;
let output_c = *output_shape.c();
let output_c_stride = *output_shape.c_stride() as isize;
let output_x_stride = output_shape.hw_strides()[0] as isize;
let output_y_stride = output_shape.hw_strides()[1] as isize;
let ox_len = output_shape.hw_dims()[0];
let oy_len = output_shape.hw_dims()[1];
let ix_len = input_shape.hw_dims()[0];
let iy_len = input_shape.hw_dims()[1];
let kx_len = self.pool_spec.kernel_shape[0];
let ky_len = self.pool_spec.kernel_shape[1];
let kvol = kx_len * ky_len;
let ihw = ix_len * iy_len;
let mut output_view = output.to_plain_array_view_mut::<T>()?;
let out_base = output_view.as_mut_ptr();
unsafe {
for n in 0..n_batch {
let out_n = out_base.add(n * *output_shape.n_stride().unwrap_or(&0));
let in_n = input.as_ptr().add(n * output_c * ihw);
for kx in 0..kx_len {
for ix in 0..ix_len {
let ox = (kx * x_dil + ix * x_stride) as isize - x_pad;
if ox < 0 || ox >= ox_len as isize {
continue;
}
let out_x = out_n.offset(ox * output_x_stride);
for ky in 0..ky_len {
let kix = kx * ky_len + ky;
let oy0 = (ky * y_dil) as isize - y_pad;
for c in 0..output_c {
let kv = *kernel.get_unchecked(c * kvol + kix);
let in_c = in_n.add(c * ihw + ix * iy_len);
let out_c = out_x.offset(c as isize * output_c_stride);
for iy in 0..iy_len {
let oy = oy0 + (iy * y_stride) as isize;
if oy < 0 || oy >= oy_len as isize {
continue;
}
*out_c.offset(oy * output_y_stride) += kv * *in_c.add(iy);
}
}
}
}
}
}
}
Ok(())
}
fn eval_t<T: Datum + Float + Copy + AddAssign<T>>(
&self,
input_shape: &DataShape,
output_shape: &DataShape,
spatial_output_details: &[ComputedPaddedDim<usize>],
kernel: &TValue,
input: &TValue,
output: &mut Tensor,
) -> TractResult<()> {
let c = *output_shape.c();
let hw: usize = input_shape.hw_dims().iter().product();
let kernel_vol: usize = self.pool_spec.kernel_shape.iter().product();
let kernel = kernel.to_plain_array_view::<T>()?;
let kernel = kernel.as_slice().context("depthwise deconv kernel must be contiguous")?;
let input = input.to_plain_array_view::<T>()?;
let input = input.as_slice().context("depthwise deconv input must be contiguous")?;
ensure!(kernel.len() == c * kernel_vol);
let n = *output_shape.n().unwrap_or(&1);
ensure!(input.len() == n * c * hw);
if input_shape.hw_rank() == 2 {
return unsafe {
self.eval_t_2d::<T>(
input_shape,
output_shape,
spatial_output_details,
kernel,
input,
output,
)
};
}
let strides = self.pool_spec.strides();
let dilations = self.pool_spec.dilations();
let mut output_plain = output.try_as_plain_ram_mut()?;
let mut output = output_plain.to_array_view_mut::<T>()?;
for n in 0..n {
for o in 0..c {
for (kix, kcoords) in
tract_ndarray::indices(&*self.pool_spec.kernel_shape).into_iter().enumerate()
{
let kv = kernel[o * kernel_vol + kix];
for (gix, gcoords) in
tract_ndarray::indices(input_shape.hw_dims()).into_iter().enumerate()
{
let ocoord: TVec<isize> = tract_itertools::izip!(
kcoords.slice(),
gcoords.slice(),
strides.as_ref(),
dilations.as_ref(),
spatial_output_details
)
.map(|(k, g, s, d, details)| {
(k * d + g * s) as isize - details.pad_before as isize
})
.collect();
if ocoord
.iter()
.zip(output_shape.hw_dims().iter())
.all(|(x, dim)| *x >= 0 && (*x as usize) < *dim)
{
let value = kv * input[(n * c + o) * hw + gix];
let ocoord = ocoord.iter().map(|x| *x as usize).collect::<TVec<_>>();
let ocoord =
self.pool_spec.data_format.with_n().from_n_c_hw(n, o, ocoord)?;
output[&*ocoord.shape] += value;
}
}
}
}
}
Ok(())
}
}
impl TypedOp for DepthwiseDeconv {
fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
ensure!(inputs.len() == 3);
let shape = super::output_shape(&self.pool_spec, &self.input_shape, &self.adjustments)?;
ensure!(*inputs[2].shape == *shape);
Ok(tvec!(inputs[1].datum_type.fact(shape)))
}
fn set_symbols(
&self,
_source: &TypedModel,
node: &TypedNode,
target: &mut TypedModel,
mapping: &HashMap<OutletId, OutletId>,
subs: &HashMap<Symbol, TDim>,
) -> TractResult<TVec<OutletId>> {
target.wire_node(
&node.name,
Self { input_shape: self.input_shape.substitute(subs)?.into_owned(), ..self.clone() },
&[mapping[&node.inputs[0]], mapping[&node.inputs[1]], mapping[&node.inputs[2]]],
)
}
as_op!();
}