Skip to content

Commit

Permalink
feat/cuda: share workspace between CUDA convolution operations
Browse files Browse the repository at this point in the history
original code by @MichaelHirn

CLOSES #18
  • Loading branch information
hobofan committed Feb 21, 2016
1 parent 62443f9 commit 7f5f320
Showing 1 changed file with 15 additions and 2 deletions.
17 changes: 15 additions & 2 deletions src/frameworks/cuda/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
)
Expand Down

0 comments on commit 7f5f320

Please sign in to comment.