Merge pull request apache#1742 from antinucleon/master

basic cudnn v5
liuyusg · Mar 30, 2016 · 68890c2 · 68890c2
2 parents a72e39a + 81988b4
commit 68890c2
Show file tree

Hide file tree

Showing 6 changed files with 81 additions and 10 deletions.
diff --git a/ps-lite b/ps-lite
diff --git a/src/operator/activation-inl.h b/src/operator/activation-inl.h
@@ -131,7 +131,7 @@ class ActivationProp : public OperatorProperty {
     CHECK_NE(dtype, -1) << "First input must have specified type";
     for (index_t i = 0; i < in_type->size(); ++i) {
       if ((*in_type)[i] == -1) {
-        (*in_type)[i] = dtype;
+          (*in_type)[i] = dtype;
       } else {
         CHECK_EQ((*in_type)[i], dtype) << "This layer requires uniform type. "
                                        << "Expected " << dtype << " v.s. given "

diff --git a/src/operator/cudnn_activation-inl.h b/src/operator/cudnn_activation-inl.h
@@ -34,11 +34,21 @@ class CuDNNActivationOp : public Operator {
         LOG(FATAL) << "Not implmented";
         break;
     }
+    #if CUDNN_MAJOR == 5
+    nan_prop_ = CUDNN_NOT_PROPAGATE_NAN;
+    CHECK_EQ(cudnnCreateActivationDescriptor(&desc_),
+             CUDNN_STATUS_SUCCESS);
+    CHECK_EQ(cudnnSetActivationDescriptor(desc_, mode_, nan_prop_, relu_ceil_),
+             CUDNN_STATUS_SUCCESS);
+    #endif
   }
 
   ~CuDNNActivationOp() {
     if (init_cudnn_) {
       CHECK_EQ(cudnnDestroyTensorDescriptor(shape_desc_), CUDNN_STATUS_SUCCESS);
+      #if CUDNN_MAJOR == 5
+      CHECK_EQ(cudnnDestroyActivationDescriptor(desc_), CUDNN_STATUS_SUCCESS);
+      #endif
     }
   }
 
@@ -89,7 +99,12 @@ class CuDNNActivationOp : public Operator {
                                           data.shape_[3]), CUDNN_STATUS_SUCCESS);
     }
     CHECK_EQ(cudnnActivationForward(s->dnn_handle_,
-                                    mode_,
+                                    #if CUDNN_MAJOR <= 4
+                                     mode_,
+                                    #endif
+                                    #if CUDNN_MAJOR == 5
+                                     desc_,
+                                    #endif
                                     &alpha,
                                     shape_desc_,
                                     data.dptr_,
@@ -145,7 +160,12 @@ class CuDNNActivationOp : public Operator {
     }
     CHECK_EQ(s->dnn_handle_ownership_, mshadow::Stream<gpu>::OwnHandle);
     CHECK_EQ(cudnnActivationBackward(s->dnn_handle_,
+                                     #if CUDNN_MAJOR <= 4
                                      mode_,
+                                     #endif
+                                     #if CUDNN_MAJOR == 5
+                                     desc_,
+                                     #endif
                                      &alpha,
                                      shape_desc_,
                                      output_data.dptr_,
@@ -164,6 +184,11 @@ class CuDNNActivationOp : public Operator {
   cudnnActivationMode_t mode_;
   cudnnTensorDescriptor_t shape_desc_;
   ActivationParam param_;
+#if CUDNN_MAJOR == 5
+  cudnnActivationDescriptor_t desc_;
+  cudnnNanPropagation_t nan_prop_;
+  double relu_ceil_;
+#endif
 };  // class CuDNNActivationOp
 }  // namespace op
 }  // namespace mxnet

diff --git a/src/operator/cudnn_convolution-inl.h b/src/operator/cudnn_convolution-inl.h
@@ -54,9 +54,8 @@ class CuDNNConvolutionOp : public Operator {
     if (!init_cudnn_) {
       Init(s, in_data, out_data);
     }
-    Tensor<gpu, 1, DType> workspace =
-        ctx.requested[conv::kTempSpace].get_space_typed<gpu, 1, DType>(
-            mshadow::Shape1(forward_workspace_), s);
+    Tensor<gpu, 1> workspace = ctx.requested[conv::kTempSpace].get_space<gpu>(
+                                 mshadow::Shape1(forward_workspace_), s);
     for (uint32_t g = 0; g < param_.num_group; ++g) {
       typename DataType<DType>::ScaleType alpha = 1.0f;
       typename DataType<DType>::ScaleType beta = 0.0f;
@@ -76,7 +75,7 @@ class CuDNNConvolutionOp : public Operator {
       if (!param_.no_bias) {
         beta = 1.0f;
         Tensor<gpu, 1, DType> bias = in_data[conv::kBias].get<gpu, 1, DType>(s);
-#if CUDNN_MAJOR == 4
+#if CUDNN_MAJOR >= 4
         CHECK_EQ(cudnnAddTensor(s->dnn_handle_,
                                 &alpha,
                                 bias_desc_,
@@ -120,8 +119,8 @@ class CuDNNConvolutionOp : public Operator {
     Tensor<gpu, 4, DType> data = in_data[conv::kData].get<gpu, 4, DType>(s);
     Tensor<gpu, 4, DType> gdata = in_grad[conv::kData].get<gpu, 4, DType>(s);
     Tensor<gpu, 1, DType> workspace =
-        ctx.requested[conv::kTempSpace].get_space_typed<gpu, 1, DType>(
-            mshadow::Shape1(backward_workspace_), s);
+      ctx.requested[conv::kTempSpace].get_space_typed<gpu, 1, DType>(
+      mshadow::Shape1(backward_workspace_), s);
     for (uint32_t g = 0; g < param_.num_group; ++g) {
       typename DataType<DType>::ScaleType alpha = 1.0f;
       typename DataType<DType>::ScaleType beta = 0.0f;
@@ -136,7 +135,12 @@ class CuDNNConvolutionOp : public Operator {
                                               gbias.dptr_ + bias_offset_ * g),
                  CUDNN_STATUS_SUCCESS);
       }
+      #if CUDNN_MAJOR <= 4
       CHECK_EQ(cudnnConvolutionBackwardFilter_v3(s->dnn_handle_,
+      #endif
+      #if CUDNN_MAJOR == 5
+      CHECK_EQ(cudnnConvolutionBackwardFilter(s->dnn_handle_,
+      #endif
                &alpha,
                in_desc_,
                data.dptr_ + data_offset_ * g,
@@ -149,7 +153,12 @@ class CuDNNConvolutionOp : public Operator {
                &beta,
                filter_desc_,
                gwmat.dptr_ + weight_offset_ * g), CUDNN_STATUS_SUCCESS);
+      #if CUDNN_MAJOR <= 4
       CHECK_EQ(cudnnConvolutionBackwardData_v3(s->dnn_handle_,
+      #endif
+      #if CUDNN_MAJOR == 5
+      CHECK_EQ(cudnnConvolutionBackwardData(s->dnn_handle_,
+      #endif
                &alpha,
                filter_desc_,
                wmat.dptr_ + weight_offset_ * g,
@@ -171,6 +180,9 @@ class CuDNNConvolutionOp : public Operator {
                    const std::vector<TBlob> &out_data) {
     using namespace mshadow;
     size_t expected = param_.no_bias ? 2 : 3;
+    #if CUDNN_MAJOR == 5
+    format_ = CUDNN_TENSOR_NCHW;
+    #endif
     CHECK_EQ(in_data.size(), expected);
     CHECK_EQ(out_data.size(), 1);
     if (!init_cudnn_) {
@@ -191,6 +203,9 @@ class CuDNNConvolutionOp : public Operator {
       CHECK_EQ(cudnnCreateConvolutionDescriptor(&conv_desc_), CUDNN_STATUS_SUCCESS);
       CHECK_EQ(cudnnSetFilter4dDescriptor(filter_desc_,
                                           dtype_,
+                                          #if CUDNN_MAJOR == 5
+                                          format_,
+                                          #endif
                                           param_.num_filter / param_.num_group,
                                           data.shape_[1] / param_.num_group,
                                           param_.kernel[0],
@@ -304,6 +319,9 @@ class CuDNNConvolutionOp : public Operator {
   cudnnConvolutionFwdAlgo_t algo_;
   cudnnConvolutionBwdDataAlgo_t back_algo_;
   cudnnConvolutionBwdFilterAlgo_t back_algo_w_;
+  #if CUDNN_MAJOR == 5
+  cudnnTensorFormat_t format_;
+  #endif
   ConvolutionParam param_;
 };
 #endif  // __CUDACC__ && CUDNN

diff --git a/src/operator/cudnn_deconvolution-inl.h b/src/operator/cudnn_deconvolution-inl.h
@@ -59,7 +59,12 @@ class CuDNNDeconvolutionOp : public Operator {
     for (uint32_t g = 0; g < param_.num_group; ++g) {
       float alpha = 1.0f;
       float beta = 0.0f;
+      #if CUDNN_MAJOR <= 4
       CHECK_EQ(cudnnConvolutionBackwardData_v3(s->dnn_handle_,
+      #endif
+      #if CUDNN_MAJOR == 5
+      CHECK_EQ(cudnnConvolutionBackwardData(s->dnn_handle_,
+      #endif
                &alpha,
                filter_desc_,
                wmat.dptr_ + weight_offset_ * g,
@@ -75,7 +80,7 @@ class CuDNNDeconvolutionOp : public Operator {
       if (!param_.no_bias) {
         beta = 1.0f;
         Tensor<gpu, 1> bias = in_data[deconv::kBias].get<gpu, 1, real_t>(s);
-#if CUDNN_MAJOR == 4
+#if CUDNN_MAJOR >= 4
         CHECK_EQ(cudnnAddTensor(s->dnn_handle_,
                                 &alpha,
                                 bias_desc_,
@@ -134,7 +139,12 @@ class CuDNNDeconvolutionOp : public Operator {
                                               gbias.dptr_ + bias_offset_ * g),
                  CUDNN_STATUS_SUCCESS);
       }
+      #if CUDNN_MAJOR <= 4
       CHECK_EQ(cudnnConvolutionBackwardFilter_v3(s->dnn_handle_,
+      #endif
+      #if CUDNN_MAJOR == 5
+      CHECK_EQ(cudnnConvolutionBackwardFilter(s->dnn_handle_,
+      #endif
                &alpha,
                out_desc_,
                grad.dptr_ + out_offset_ * g,
@@ -168,6 +178,9 @@ class CuDNNDeconvolutionOp : public Operator {
                    const std::vector<TBlob> &in_data,
                    const std::vector<TBlob> &out_data) {
     using namespace mshadow;
+    #if CUDNN_MAJOR == 5
+    format_ = CUDNN_TENSOR_NCHW;
+    #endif
     size_t expected = param_.no_bias ? 2 : 3;
     CHECK_EQ(in_data.size(), expected);
     CHECK_EQ(out_data.size(), 1);
@@ -189,6 +202,9 @@ class CuDNNDeconvolutionOp : public Operator {
       CHECK_EQ(cudnnCreateConvolutionDescriptor(&conv_desc_), CUDNN_STATUS_SUCCESS);
       CHECK_EQ(cudnnSetFilter4dDescriptor(filter_desc_,
                                           dtype_,
+                                          #if CUDNN_MAJOR == 5
+                                          format_,
+                                          #endif
                                           data.shape_[1] / param_.num_group,
                                           param_.num_filter / param_.num_group,
                                           param_.kernel[0],
@@ -302,6 +318,9 @@ class CuDNNDeconvolutionOp : public Operator {
   cudnnConvolutionFwdAlgo_t algo_;
   cudnnConvolutionBwdDataAlgo_t back_algo_;
   cudnnConvolutionBwdFilterAlgo_t back_algo_w_;
+  #if CUDNN_MAJOR == 5
+  cudnnTensorFormat_t format_;
+  #endif
   DeconvolutionParam param_;
 };
 #endif  // __CUDACC__ && CUDNN

diff --git a/src/operator/cudnn_pooling-inl.h b/src/operator/cudnn_pooling-inl.h
@@ -113,6 +113,9 @@ class CuDNNPoolingOp : public Operator {
                    const std::vector<TBlob> &in_data,
                    const std::vector<TBlob> &out_data) {
     using namespace mshadow;
+    #if CUDNN_MAJOR == 5
+    nan_prop_ = CUDNN_NOT_PROPAGATE_NAN;
+    #endif
     CHECK_EQ(in_data.size(), 1);
     CHECK_EQ(out_data.size(), 1);
     if (!init_cudnn_) {
@@ -138,6 +141,9 @@ class CuDNNPoolingOp : public Operator {
                                           out.shape_[3]), CUDNN_STATUS_SUCCESS);
       CHECK_EQ(cudnnSetPooling2dDescriptor(pooling_desc_,
                                            mode_,
+                                           #if CUDNN_MAJOR == 5
+                                           nan_prop_,
+                                           #endif
                                            param_.kernel[0],
                                            param_.kernel[1],
                                            param_.pad[0],
@@ -153,6 +159,9 @@ class CuDNNPoolingOp : public Operator {
   cudnnTensorDescriptor_t in_desc_;
   cudnnTensorDescriptor_t out_desc_;
   cudnnPoolingDescriptor_t pooling_desc_;
+  #if CUDNN_MAJOR == 5
+  cudnnNanPropagation_t nan_prop_;
+  #endif
   PoolingParam param_;
 };  // class CuDNNPoolingOp
 }  // namespace op
+6 −2		.gitignore
+4 −0		CMakeLists.txt
+2 −2		Makefile
+4 −4		docs/Doxyfile
+216 −0		docs/Makefile
+40 −0		docs/api.md
+315 −0		docs/conf.py
+0 −0		docs/env.md
+1 −0		docs/get_started.md
+58 −0		docs/how_to.md
+16 −0		docs/index.md
+184 −0		docs/overview.md
+1 −0		docs/requirements.txt
+19 −0		docs/sphinx_util.py
+1 −0		docs/tutorials.md
+0 −5		include/ps/base.h
+62 −12		include/ps/internal/message.h
+6 −29		include/ps/internal/postoffice.h
+0 −25		include/ps/internal/van.h
+0 −9		include/ps/ps.h
+1 −0		include/ps/sarray.h
+27 −37		make/deps.mk
+1 −1		src/meta.proto
+4 −8		src/network_utils.h
+1 −0		src/postoffice.cc
+70 −45		src/van.cc
+3 −0		tracker/dmlc_mpi.py