Add Cuda RNN FFI

cudnn-sys/build.rs

@@ -62,7 +62,7 @@ fn main() {
             // requires a nightly rustc and enabling
             // unstable features.
             .rust_target(bindgen::RustTarget::Stable_1_19)
-            .hide_type("max_align_t") // https://github.com/servo/rust-bindgen/issues/550
+            .blacklist_type("max_align_t") // https://github.com/servo/rust-bindgen/issues/550
             .raw_line(
                 r"
 //! Defines the FFI for CUDA cuDNN.

cudnn/src/api/utils.rs

@@ -3,6 +3,7 @@
 use crate::ffi::*;
 use std::ptr;
 use crate::{Error, API};
+use utils::DataType;
 
 impl API {
     /// Initialize the CUDA cuDNN API with needed context and resources.
@@ -53,3 +54,15 @@ impl API {
         }
     }
 }
+
+impl API {
+    /// Convert to CUDNN Data Type
+    pub fn to_cudnn_data_type(data_type: DataType) -> cudnnDataType_t {
+        let data_type = match data_type {
+            DataType::Float => cudnnDataType_t::CUDNN_DATA_FLOAT,
+            DataType::Double => cudnnDataType_t::CUDNN_DATA_DOUBLE,
+            DataType::Half => cudnnDataType_t::CUDNN_DATA_HALF
+        };
+        data_type
+    }
+}

cudnn/src/cuda.rs

@@ -38,8 +38,8 @@ impl CudaDeviceMemory {
     }
 
     /// Returns the size of the CUDA Device Memory chunk.
-    pub fn size(&self) -> &usize {
-        &self.size
+    pub fn size(&self) -> usize {
+        self.size
     }
 }
 

cudnn/src/cudnn.rs

@@ -6,13 +6,15 @@
 
 use super::utils::{
     ActivationConfig, ConvolutionConfig, DataTypeInfo, DropoutConfig, NormalizationConfig,
-    PoolingConfig, ScalParams,
+    PoolingConfig, ScalParams, RnnConfig
 };
 use super::*;
 
 use crate::cuda::CudaDeviceMemory;
 use num::traits::Float;
 use std::mem::transmute_copy;
+use utils::DataType;
+use tensor_descriptor::tensor_vec_id_c;
 
 #[derive(Debug, Clone)]
 /// Provides a the high-level interface to CUDA's cuDNN.
@@ -178,6 +180,116 @@ impl Cudnn {
         Ok(DropoutConfig::new(dropout, reserve))
     }
 
+    /// Initialize RNN
+    pub fn init_rnn(
+        &self,
+        x_desc: &[TensorDescriptor],
+        rnn_desc: RnnDescriptor,
+        hidden_size: i32,
+        num_layers: i32,
+        seq_length: i32,
+        dropout_desc: &DropoutDescriptor,
+        input_mode: cudnnRNNInputMode_t,
+        direction_mode: cudnnDirectionMode_t,
+        network_mode: cudnnRNNMode_t,
+        algorithm: cudnnRNNAlgo_t,
+        data_type: DataType,
+        math_type: cudnnMathType_t
+    ) -> Result<RnnConfig, Error> {
+
+        let data_type =  match data_type {
+            DataType::Float => cudnnDataType_t::CUDNN_DATA_FLOAT,
+            DataType::Double => cudnnDataType_t::CUDNN_DATA_DOUBLE,
+            DataType::Half => cudnnDataType_t::CUDNN_DATA_HALF
+        };
+
+        API::set_rnn_matrix_math_type(
+            *rnn_desc.id_c(),
+            math_type
+        )?;
+
+        let workspace_size : usize = API::get_rnn_workspace_size(
+            *self.id_c(),
+            *rnn_desc.id_c(),
+            seq_length,
+            tensor_vec_id_c(x_desc)
+        )?;
+
+        let training_reserve_size : usize = API::get_rnn_training_reserve_size(
+            *self.id_c(),
+            *rnn_desc.id_c(),
+            seq_length,
+            tensor_vec_id_c(x_desc)
+        )?;
+
+        let training_reserve : CudaDeviceMemory = CudaDeviceMemory::new(training_reserve_size)?;
+
+        Ok(RnnConfig::new(
+            rnn_desc,
+            hidden_size,
+            num_layers,
+            seq_length,
+            *dropout_desc.id_c(),
+            input_mode,
+            direction_mode,
+            network_mode,
+            algorithm,
+            data_type,
+            workspace_size,
+            training_reserve_size,
+            training_reserve
+        ))
+    }
+
+    /// Train & Return Results for RNN
+    pub fn rnn_forward<T>(
+        &self,
+        rnn_config: &RnnConfig,
+        src_desc: Vec<TensorDescriptor>,
+        src: *const ::libc::c_void,
+        output_desc: Vec<TensorDescriptor>,
+        output: *mut ::libc::c_void,
+        hidden_desc: &TensorDescriptor,
+        // Planning to initially pass NULLs to this
+        hidden: *const ::libc::c_void,
+        cell_desc: &TensorDescriptor,
+        // Planning to initially pass NULLs to this
+        cell: *const ::libc::c_void,
+        weight_desc: &FilterDescriptor,
+        weight: *const ::libc::c_void,
+        hidden_output_desc: &TensorDescriptor,
+        hidden_output: *mut ::libc::c_void,
+        cell_output_desc: &TensorDescriptor,
+        cell_output: *mut ::libc::c_void,
+        workspace: *mut ::libc::c_void,
+        reserve_data: *mut ::libc::c_void
+    ) -> Result<(), Error>
+    where T: Float + DataTypeInfo {
+        API::rnn_forward_training(
+            *self.id_c(),
+            *(rnn_config.rnn_desc().id_c()),
+            *rnn_config.sequence_length(),
+            tensor_vec_id_c(&src_desc),
+            src,
+            *hidden_desc.id_c(),
+            hidden,
+            *cell_desc.id_c(),
+            cell,
+            *weight_desc.id_c(),
+            weight,
+            tensor_vec_id_c(&output_desc),
+            output,
+            *hidden_output_desc.id_c(),
+            hidden_output,
+            *cell_output_desc.id_c(),
+            cell_output,
+            workspace,
+            rnn_config.rnn_workspace_size(),
+            reserve_data,
+            rnn_config.training_reserve_size()
+        )
+    }
+
     /// Computes the forward Sigmoid Activation function.
     ///
     /// Writes the result of the computation to `dest_data`.
@@ -386,7 +498,7 @@ impl Cudnn {
             *conv_config.forward_algo(),
             *conv_config.conv_desc().id_c(),
             workspace,
-            *conv_config.forward_workspace_size(),
+            conv_config.forward_workspace_size(),
             unsafe { transmute_copy(&&scale.a) },
             *src_desc.id_c(),
             src_data,
@@ -445,7 +557,7 @@ impl Cudnn {
             *conv_config.backward_filter_algo(),
             *conv_config.conv_desc().id_c(),
             workspace,
-            *conv_config.backward_filter_workspace_size(),
+            conv_config.backward_filter_workspace_size(),
             unsafe { transmute_copy(&&scale.a) },
             *src_desc.id_c(),
             src_data,
@@ -479,7 +591,7 @@ impl Cudnn {
             *conv_config.backward_data_algo(),
             *conv_config.conv_desc().id_c(),
             workspace,
-            *conv_config.backward_data_workspace_size(),
+            conv_config.backward_data_workspace_size(),
             unsafe { transmute_copy(&&scale.a) },
             *conv_config.filter_desc().id_c(),
             filter_data,
@@ -817,7 +929,7 @@ impl Cudnn {
             *dest_desc.id_c(),
             dest_data,
             *dropout_conf.reserved_space().id_c(),
-            *dropout_conf.reserved_space().size(),
+            dropout_conf.reserved_space().size(),
         )
     }
 
@@ -843,7 +955,7 @@ impl Cudnn {
             *dest_desc.id_c(),
             dest_data,
             *dropout_conf.reserved_space().id_c(),
-            *dropout_conf.reserved_space().size(),
+            dropout_conf.reserved_space().size(),
         )
     }
 }

cudnn/src/dropout_descriptor.rs

@@ -36,7 +36,7 @@ impl DropoutDescriptor {
             *handle.id_c(),
             dropout,
             *reserve.id_c(),
-            *reserve.size(),
+            reserve.size(),
             seed,
         )?;
 

cudnn/src/lib.rs

@@ -83,6 +83,7 @@ pub use self::filter_descriptor::FilterDescriptor;
 pub use self::normalization_descriptor::NormalizationDescriptor;
 pub use self::pooling_descriptor::PoolingDescriptor;
 pub use self::tensor_descriptor::TensorDescriptor;
+pub use self::rnn_descriptor::RnnDescriptor;
 pub use crate::ffi::*;
 
 #[derive(Debug, Copy, Clone)]
@@ -100,4 +101,5 @@ mod filter_descriptor;
 mod normalization_descriptor;
 mod pooling_descriptor;
 mod tensor_descriptor;
+mod rnn_descriptor;
 pub mod utils;

cudnn/src/rnn_descriptor.rs

@@ -6,11 +6,12 @@
 use super::{API, Error};
 use super::utils::DataType;
 use ffi::*;
+use ::{Cudnn, DropoutDescriptor};
 
 #[derive(Debug, Clone)]
 /// Describes a Recurrent Descriptor.
 pub struct RnnDescriptor {
-    id: cudnnRnnDescriptor_t,
+    id: cudnnRNNDescriptor_t,
 }
 
 impl Drop for RnnDescriptor {
@@ -26,48 +27,39 @@ impl RnnDescriptor {
         handle: &Cudnn,
         hidden_size: i32,
         num_layers: i32,
-        dropout_desc: DropoutDescriptor,
+        dropout_desc: &DropoutDescriptor,
         input_mode: cudnnRNNInputMode_t,
         direction: cudnnDirectionMode_t,
         mode: cudnnRNNMode_t,
         algorithm: cudnnRNNAlgo_t,
         data_type: DataType,
     ) -> Result<RnnDescriptor, Error> {
-        let array_length = pad.len() as i32;
-        let upscale: Vec<i32> = ::std::iter::repeat(1i32)
-            .take(array_length as usize)
-            .collect();
-
         let generic_rnn_desc = API::create_rnn_descriptor()?;
-        let data_type = match data_type {
-            DataType::Float => cudnnDataType_t::CUDNN_DATA_FLOAT,
-            DataType::Double => cudnnDataType_t::CUDNN_DATA_DOUBLE,
-            DataType::Half => cudnnDataType_t::CUDNN_DATA_HALF,
-            _ => return Err(Error::InvalidValue("Invalid data type value passed")),
-        };
-        API::set_rnn_descriptor(
-            handle.id_c(),
+         API::set_rnn_descriptor(
+            *handle.id_c(),
             generic_rnn_desc,
             hidden_size,
             num_layers,
-            dropout_desc,
+            *dropout_desc.id_c(),
             input_mode,
             direction,
             mode,
             algorithm,
             data_type,
-        );
-        Ok(RnnDescriptor::from_c(generic_rnn_desc))
+        )?;
 
+        Ok(RnnDescriptor {
+            id: generic_rnn_desc
+        })
     }
 
     /// Initializes a new CUDA cuDNN RnnDescriptor from its C type.
-    pub fn from_c(id: cudnnRnnDescriptor_t) -> RnnDescriptor {
+    pub fn from_c(id: cudnnRNNDescriptor_t) -> RnnDescriptor {
         RnnDescriptor { id: id }
     }
 
     /// Returns the CUDA cuDNN RnnDescriptor as its C type.
-    pub fn id_c(&self) -> &cudnnRnnDescriptor_t {
+    pub fn id_c(&self) -> &cudnnRNNDescriptor_t {
         &self.id
     }
 }

cudnn/src/tensor_descriptor.rs

@@ -14,6 +14,13 @@ pub struct TensorDescriptor {
     id: cudnnTensorDescriptor_t,
 }
 
+/// Return C Handle for a Vector of Tensor Descriptors
+pub fn tensor_vec_id_c(tensor_vec: &[TensorDescriptor]) -> Vec<cudnnTensorDescriptor_t> {
+    tensor_vec.iter().map(|tensor| {
+        *(*tensor).id_c()
+    }).collect()
+}
+
 impl Drop for TensorDescriptor {
     #[allow(unused_must_use)]
     fn drop(&mut self) {
@@ -85,3 +92,4 @@ impl TensorDescriptor {
         &self.id
     }
 }
+