Implement 3D deconvolution (cuDNN)

src/operator/cudnn_algoreg-inl.h

@@ -0,0 +1,89 @@
+/*!
+ * Copyright (c) 2015 by Contributors
+ * \file cudnn_algoreg-inl.h
+ * \brief
+ * \author Bing Xu
+ */
+#ifndef MXNET_OPERATOR_CUDNN_ALGOREG_INL_H_
+#define MXNET_OPERATOR_CUDNN_ALGOREG_INL_H_
+
+#include <algorithm>
+#include <mutex>
+#include <string>
+#include <vector>
+#include "../common/cuda_utils.h"
+#include "./convolution-inl.h"
+#include "./deconvolution-inl.h"
+
+namespace mxnet {
+namespace op {
+#if MXNET_USE_CUDNN == 1
+
+class CuDNNAlgoReg {
+ public:
+  template <typename Param>
+  std::string GetKey(const Param &param, const std::vector<TShape> &in_shape,
+                     const std::vector<TShape> &out_shape) {
+    std::ostringstream oss;
+    for (auto &i : in_shape)
+      oss << i << ";";
+    for (auto &i : out_shape)
+      oss << i << ";";
+    auto dict = param.__DICT__();
+    for (auto &k : dict)
+      oss << k.first << "=" << k.second << ";";
+    return oss.str();
+  }
+
+  bool Find(std::string key, cudnnConvolutionFwdAlgo_t *fwd,
+            cudnnConvolutionBwdDataAlgo_t *bwd,
+            cudnnConvolutionBwdFilterAlgo_t *flt) {
+    std::lock_guard<std::mutex> guard(lock_);
+    auto i = reg_.find(key);
+    if (i != reg_.end()) {
+      *fwd = i->second.fwd;
+      *bwd = i->second.bwd;
+      *flt = i->second.flt;
+      return true;
+    }
+    return false;
+  }
+
+  void Register(std::string key, cudnnConvolutionFwdAlgo_t fwd,
+                cudnnConvolutionBwdDataAlgo_t bwd,
+                cudnnConvolutionBwdFilterAlgo_t flt) {
+    std::lock_guard<std::mutex> guard(lock_);
+    if (reg_.size() % 50 == 0) {
+      LOG(INFO) << "Running performance tests to find the best convolution "
+                   "algorithm, "
+                   "this can take a while... (setting env variable "
+                   "MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable)";
+      if (reg_.size() >= 1000) {
+        LOG(INFO)
+            << "If you see this message in the middle of training, you are "
+               "probably using bucketing. Consider setting env variable "
+               "MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable cudnn tuning.";
+      }
+    }
+    reg_[key].fwd = fwd;
+    reg_[key].bwd = bwd;
+    reg_[key].flt = flt;
+  }
+
+  static CuDNNAlgoReg *Get();
+
+ private:
+  struct CudnnAlgorithms {
+    cudnnConvolutionFwdAlgo_t fwd;
+    cudnnConvolutionBwdDataAlgo_t bwd;
+    cudnnConvolutionBwdFilterAlgo_t flt;
+  };
+
+  std::mutex lock_;
+  std::unordered_map<std::string, CudnnAlgorithms> reg_;
+};
+#endif  // __CUDACC__ && CUDNN
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // MXNET_OPERATOR_CUDNN_ALGOREG_INL_H_

src/operator/cudnn_convolution.cc → src/operator/cudnn_algoreg.cc

@@ -1,10 +1,10 @@
 /*!
  * Copyright (c) 2015 by Contributors
- * \file cudnn_convolution.cc
+ * \file cudnn_algoreg.cc
  * \brief
  * \author Junyuan Xie
 */
-#include "./cudnn_convolution-inl.h"
+#include "./cudnn_algoreg-inl.h"
 #include <mxnet/base.h>
 #include <mxnet/ndarray.h>
 
@@ -14,8 +14,8 @@
 namespace mxnet {
 namespace op {
 #if MXNET_USE_CUDNN == 1
-CuDNNAlgoReg* CuDNNAlgoReg::Get() {
-  static CuDNNAlgoReg* ptr = new CuDNNAlgoReg();
+CuDNNAlgoReg *CuDNNAlgoReg::Get() {
+  static CuDNNAlgoReg *ptr = new CuDNNAlgoReg();
   return ptr;
 }
 #endif  // CUDNN

src/operator/cudnn_convolution-inl.h

@@ -12,75 +12,13 @@
 #include <mutex>
 #include <string>
 #include "./convolution-inl.h"
+#include "./cudnn_algoreg-inl.h"
 #include "../common/cuda_utils.h"
 
 namespace mxnet {
 namespace op {
 #if MXNET_USE_CUDNN == 1
 
-class CuDNNAlgoReg {
- public:
-  std::string GetKey(const ConvolutionParam& param,
-                     const std::vector<TShape>& in_shape,
-                     const std::vector<TShape>& out_shape) {
-    std::ostringstream oss;
-    for (auto& i : in_shape) oss << i << ";";
-    for (auto& i : out_shape) oss << i << ";";
-    auto dict = param.__DICT__();
-    for (auto& k : dict) oss << k.first << "=" << k.second << ";";
-    return oss.str();
-  }
-
-  bool Find(std::string key,
-            cudnnConvolutionFwdAlgo_t *fwd,
-            cudnnConvolutionBwdDataAlgo_t *bwd,
-            cudnnConvolutionBwdFilterAlgo_t *flt) {
-    std::lock_guard<std::mutex> guard(lock_);
-    auto i = reg_.find(key);
-    if (i != reg_.end()) {
-      *fwd = i->second.fwd;
-      *bwd = i->second.bwd;
-      *flt = i->second.flt;
-      return true;
-    }
-    return false;
-  }
-
-  void Register(std::string key,
-                cudnnConvolutionFwdAlgo_t fwd,
-                cudnnConvolutionBwdDataAlgo_t bwd,
-                cudnnConvolutionBwdFilterAlgo_t flt) {
-    std::lock_guard<std::mutex> guard(lock_);
-    if (reg_.size() % 50 == 0) {
-      LOG(INFO)
-        << "Running performance tests to find the best convolution algorithm, "
-           "this can take a while... (setting env variable "
-           "MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable)";
-      if (reg_.size() >= 1000) {
-        LOG(INFO)
-          << "If you see this message in the middle of training, you are "
-             "probably using bucketing. Consider setting env variable "
-             "MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable cudnn tuning.";
-      }
-    }
-    reg_[key].fwd = fwd;
-    reg_[key].bwd = bwd;
-    reg_[key].flt = flt;
-  }
-
-  static CuDNNAlgoReg* Get();
-
- private:
-  struct CudnnAlgorithms {
-    cudnnConvolutionFwdAlgo_t fwd;
-    cudnnConvolutionBwdDataAlgo_t bwd;
-    cudnnConvolutionBwdFilterAlgo_t flt;
-  };
-
-  std::mutex lock_;
-  std::unordered_map<std::string, CudnnAlgorithms> reg_;
-};
-
 template<typename DType>
 class CuDNNConvolutionOp : public Operator {
  public: