From eb2b848df173f7a07eb0d76a432c5d4badca7ba6 Mon Sep 17 00:00:00 2001
From: Felix Abecassis <fabecassis@nvidia.com>
Date: Thu, 10 Dec 2015 15:11:51 -0800
Subject: [PATCH] Fix CuDNNConvolutionLayer for cuDNN v4

Add a macro to check the current cuDNN version
---
 include/caffe/util/cudnn.hpp         | 3 +++
 src/caffe/layers/cudnn_conv_layer.cu | 8 ++++++++
 2 files changed, 11 insertions(+)

diff --git a/include/caffe/util/cudnn.hpp b/include/caffe/util/cudnn.hpp
index b531dd5fa7a..8a7e17c6cd4 100644
--- a/include/caffe/util/cudnn.hpp
+++ b/include/caffe/util/cudnn.hpp
@@ -7,6 +7,9 @@
 #include "caffe/common.hpp"
 #include "caffe/proto/caffe.pb.h"
 
+#define CUDNN_VERSION_MIN(major, minor, patch) \
+    (CUDNN_VERSION >= (major * 1000 + minor * 100 + patch))
+
 #define CUDNN_CHECK(condition) \
   do { \
     cudnnStatus_t status = condition; \
diff --git a/src/caffe/layers/cudnn_conv_layer.cu b/src/caffe/layers/cudnn_conv_layer.cu
index 1990e932a70..42c4fd0260c 100644
--- a/src/caffe/layers/cudnn_conv_layer.cu
+++ b/src/caffe/layers/cudnn_conv_layer.cu
@@ -30,11 +30,19 @@ void CuDNNConvolutionLayer<Dtype>::Forward_gpu(
       // Bias.
       if (this->bias_term_) {
         const Dtype* bias_data = this->blobs_[1]->gpu_data();
+#if CUDNN_VERSION_MIN(4, 0, 0)
+        CUDNN_CHECK(cudnnAddTensor(handle_[g],
+              cudnn::dataType<Dtype>::one,
+              bias_desc_, bias_data + bias_offset_ * g,
+              cudnn::dataType<Dtype>::one,
+              top_descs_[i], top_data + top_offset_ * g));
+#else
         CUDNN_CHECK(cudnnAddTensor(handle_[g], CUDNN_ADD_SAME_C,
               cudnn::dataType<Dtype>::one,
               bias_desc_, bias_data + bias_offset_ * g,
               cudnn::dataType<Dtype>::one,
               top_descs_[i], top_data + top_offset_ * g));
+#endif
       }
     }