From 7f5f3207873874accb7a5a16d637e2701161ac04 Mon Sep 17 00:00:00 2001 From: Hobofan Date: Sun, 21 Feb 2016 19:21:56 +0100 Subject: [PATCH] feat/cuda: share workspace between CUDA convolution operations original code by @MichaelHirn CLOSES #18 --- src/frameworks/cuda/mod.rs | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/frameworks/cuda/mod.rs b/src/frameworks/cuda/mod.rs index 878fb51..490dfdf 100644 --- a/src/frameworks/cuda/mod.rs +++ b/src/frameworks/cuda/mod.rs @@ -321,11 +321,24 @@ macro_rules! impl_convolution_for_cuda_backend { } }; + // share one workspace to reduce memory + let workspace: ::co::frameworks::cuda::Memory; + if workspace_size_bwd_data >= workspace_size_bwd_filter && workspace_size_bwd_data >= workspace_size_fwd { + workspace = workspace_bwd_data; + } else if workspace_size_bwd_filter >= workspace_size_bwd_data && workspace_size_bwd_filter >= workspace_size_fwd { + workspace = workspace_bwd_filter; + } else { + workspace = workspace_fwd; + } + + let workspace_bwd_fiter = ::co::frameworks::cuda::Memory::from_c(*workspace.id_c()); + let workspace_fwd = ::co::frameworks::cuda::Memory::from_c(*workspace.id_c()); + Ok( ::cudnn::utils::ConvolutionConfig::new( useable_algo_fwd.as_cudnn().unwrap(), workspace_fwd, workspace_size_fwd, - useable_algo_bwd_filter.as_cudnn().unwrap(), workspace_bwd_filter, workspace_size_bwd_filter, - useable_algo_bwd_data.as_cudnn().unwrap(), workspace_bwd_data, workspace_size_bwd_data, + useable_algo_bwd_filter.as_cudnn().unwrap(), workspace_bwd_fiter, workspace_size_bwd_filter, + useable_algo_bwd_data.as_cudnn().unwrap(), workspace, workspace_size_bwd_data, conv_desc, filter_desc ) )