From 36f305bb0baeeafaaf85d93d1a29d09fccedb473 Mon Sep 17 00:00:00 2001
From: Max Freudenberg <67329240+maxfreu@users.noreply.github.com>
Date: Wed, 15 Jun 2022 11:09:47 +0200
Subject: [PATCH] save allocs during algorithm search

Sets allocateTmpBuf in `cudnnConvolutionXXAlgoPerf` to false if beta is zero to save an allocation in this case.
---
 ext/CUDAExt/src/cudnn/conv.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ext/CUDAExt/src/cudnn/conv.jl b/ext/CUDAExt/src/cudnn/conv.jl
index 712a2f460..e25796078 100644
--- 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