save allocs during algorithm search

Sets allocateTmpBuf in `cudnnConvolutionXXAlgoPerf` to false if beta is zero to save an allocation in this case.
FluxML · Jun 16, 2022 · 36f305b · 36f305b
1 parent bc07266
commit 36f305b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ext/CUDAExt/src/cudnn/conv.jl b/ext/CUDAExt/src/cudnn/conv.jl
@@ -97,7 +97,7 @@ function ∇conv_data!(dx::DenseCuArray{T}, dy::DenseCuArray{T}, w::DenseCuArray
     alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta);
     convDesc, dx, depad = cudnnConvolutionDescriptorAndPaddedInput(cdims, dx)
     xDesc, yDesc, wDesc = cudnnTensorDescriptor(dx), cudnnTensorDescriptor(dy), cudnnFilterDescriptor(w)
-    p = cudnnConvolutionBwdDataAlgoPerf(wDesc, w, yDesc, dy, convDesc, xDesc, dx)
+    p = cudnnConvolutionBwdDataAlgoPerf(wDesc, w, yDesc, dy, convDesc, xDesc, dx, beta!=0)
     with_workspace(p.memory) do workspace
         cudnnConvolutionBackwardData(handle(), alpha, wDesc, w, yDesc, dy, convDesc, p.algo, workspace, sizeof(workspace), beta, xDesc, dx)
     end
@@ -115,7 +115,7 @@ function ∇conv_filter!(dw::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArr
     alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta);
     convDesc, x, _ = cudnnConvolutionDescriptorAndPaddedInput(cdims, x)
     xDesc, yDesc, wDesc = cudnnTensorDescriptor(x), cudnnTensorDescriptor(dy), cudnnFilterDescriptor(dw)
-    p = cudnnConvolutionBwdFilterAlgoPerf(xDesc, x, yDesc, dy, convDesc, wDesc, dw);
+    p = cudnnConvolutionBwdFilterAlgoPerf(xDesc, x, yDesc, dy, convDesc, wDesc, dw, beta!=0);
     with_workspace(p.memory) do workspace
         cudnnConvolutionBackwardFilter(handle(), alpha, xDesc, x, yDesc, dy, convDesc, p.algo, workspace, sizeof(workspace), beta, wDesc, dw);
     end