Add basic convolution functionalities for forward and backward passes (#29)

Author: frjnn

crates/cudnn/src/context.rs

@@ -1,18 +1,20 @@
 use crate::{
-    convolution_mode::ConvolutionMode,
+    convolution::ConvolutionMode,
     data_type::DataType,
     error::{CudnnError, IntoResult},
     math_type::MathType,
     sys,
-    tensor_format::TensorFormat,
+    tensor::TensorFormat,
 };
 
 use std::{marker::PhantomData, mem::MaybeUninit};
 
 /// A generic description of an n-dimensional convolution.
 ///
-/// *Do note* that N can be either 2 or 3, respectively for a 2-d or a 3-d convolution.
-#[derive(Debug, Clone, PartialEq, Hash)]
+/// **Do note** that N can be either 2 or 3, respectively for a 2-d or a 3-d convolution, and that
+/// the same convolution descriptor can be reused in the backward path provided it corresponds to
+/// the same layer.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct ConvolutionDescriptor<T: DataType, const N: usize> {
     pub(crate) raw: sys::cudnnConvolutionDescriptor_t,
     comp_type: PhantomData<T>,
@@ -131,6 +133,12 @@ impl<T: DataType, const N: usize> ConvolutionDescriptor<T, N> {
     }
 
     /// Sets the `MathType` for this convolution descriptor instance.
+    ///
+    /// # Arguments
+    ///
+    /// `math_type` - the provided math type.
+    ///
+    /// **Do note** that tensor core operations may not be available on all device architectures.
     pub fn set_math_type(&mut self, math_type: MathType) -> Result<(), CudnnError> {
         unsafe { sys::cudnnSetConvolutionMathType(self.raw, math_type.into()).into_result() }
     }

crates/cudnn/src/convolution/convolution_algo.rs

@@ -1,9 +1,11 @@
 use crate::sys;
 
-/// Enum used to configure a convolution descriptor. The filter used for the convolution can be
-/// applied in two different ways, corresponding mathematically to a convolution or to a
-/// cross-correlation. A cross-correlation is equivalent to a convolution with its filter
-/// rotated by 180 degrees.
+/// Enum used to configure a convolution descriptor.
+///
+/// The filter used for the convolution can be applied in two different ways, corresponding
+/// mathematically to a convolution or to a cross-correlation.
+///
+/// A cross-correlation is equivalent to a convolution with its filter rotated by 180 degrees.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum ConvolutionMode {
     /// Convolution operation.
@@ -22,7 +24,6 @@ impl From<sys::cudnnConvolutionMode_t> for ConvolutionMode {
 }
 
 impl From<ConvolutionMode> for sys::cudnnConvolutionMode_t {
-    /// Returns the raw cuDNN type associated to the given variant.
     fn from(convolution_mode: ConvolutionMode) -> sys::cudnnConvolutionMode_t {
         match convolution_mode {
             ConvolutionMode::Convolution => sys::cudnnConvolutionMode_t::CUDNN_CONVOLUTION,

crates/cudnn/src/convolution_descriptor.rs renamed to crates/cudnn/src/convolution/convolution_descriptor.rs

@@ -2,7 +2,7 @@ use crate::{
     data_type::DataType,
     error::{CudnnError, IntoResult},
     sys,
-    tensor_format::{SupportedType, TensorFormat},
+    tensor::{SupportedType, TensorFormat},
 };
 use std::{
     marker::PhantomData,

crates/cudnn/src/convolution_mode.rs renamed to crates/cudnn/src/convolution/convolution_mode.rs

@@ -1,8 +1,10 @@
+mod filter_descriptor;
+pub use filter_descriptor::*;
+
 use crate::{
     data_type::DataType,
     error::CudnnError,
-    filter_descriptor::FilterDescriptor,
-    tensor_format::{SupportedType, TensorFormat},
+    tensor::{SupportedType, TensorFormat},
 };
 use cust::memory::{DeviceBuffer, DeviceCopy, GpuBox, GpuBuffer, UnifiedBuffer};
 

crates/cudnn/src/filter_descriptor.rs renamed to crates/cudnn/src/convolution/filter/filter_descriptor.rs

@@ -0,0 +1,9 @@
+mod convolution_algo;
+mod convolution_descriptor;
+mod convolution_mode;
+mod filter;
+
+pub use convolution_algo::*;
+pub use convolution_descriptor::*;
+pub use convolution_mode::*;
+pub use filter::*;