ruDNN 0.1.7

Neural network operators for Ruda.
Documentation
use {ruda_kernel::dsl::Runtime, ruda_kernel::tensor::RudaTensor};
use {ruda_core::tensor::spatial::ConvTransposeOptions};
use {crate::convolution::components::ConvSetupError};

#[cfg(feature = "tensor-convolution-autotune")]
use {super::conv_transpose2d_autotune};
use {super::conv_transpose2d_col2im, super::conv_transpose2d_direct};

/// The strategy to be used when launching a conv_transpose kernel.
pub enum ConvTranspose2dStrategy {
    /// A simple direct convolution.
    Direct,
    #[cfg(feature = "tensor-convolution-autotune")]
    /// Using autotune to choose the best kernel based on runtime information.
    Autotune,
    /// GEMM (im2col) based implementation of convolution. Significantly increased memory usage.
    Gemm,
}

impl Default for ConvTranspose2dStrategy {
    fn default() -> Self {
        // if autotune is enabled, default to autotune
        #[cfg(feature = "tensor-convolution-autotune")]
        return ConvTranspose2dStrategy::Autotune;

        // if autotune is disabled, default to the more memory-conservative algorithm
        #[cfg(not(feature = "tensor-convolution-autotune"))]
        ConvTranspose2dStrategy::Direct
    }
}

/// Performs a 2D convolution with the given strategy
///
/// * `input` - The input feature map
/// * `weight` - The weights (filter) applied to each kernel
/// * `bias` - The bias added to each channel
/// * `options` - The options to use for the convolution
/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option.
pub fn conv_transpose2d<R: Runtime>(
    input: RudaTensor<R>,
    weight: RudaTensor<R>,
    bias: Option<RudaTensor<R>>,
    options: ConvTransposeOptions<2>,
    strategy: ConvTranspose2dStrategy,
) -> Result<RudaTensor<R>, ConvSetupError> {
    match strategy {
        ConvTranspose2dStrategy::Direct => conv_transpose2d_direct(input, weight, bias, options),
        #[cfg(feature = "tensor-convolution-autotune")]
        ConvTranspose2dStrategy::Autotune => {
            Ok(conv_transpose2d_autotune(input, weight, bias, options))
        }
        ConvTranspose2dStrategy::Gemm => conv_transpose2d_col2im(input, weight, bias, options),
    }
}