add cudnn back

7 files changed, 209 insertions, 210 deletions

diff --git a/src/examples/mnist/model_lenet.h b/src/examples/mnist/model_lenet.h
index 8dd96634..968ceaf3 100644
--- a/src/examples/mnist/model_lenet.h
+++ b/src/examples/mnist/model_lenet.h
@@ -33,9 +33,6 @@ protected:
 
     // Construct hidden layers
 
-    ABORT("TEMPORARY");
-
-    /*
     auto conv_1 = convolution(g)
                     ("prefix", "conv_1")
                     ("kernel-dims", std::make_pair(3,3))
@@ -102,7 +99,6 @@ protected:
       // Define a top-level node for inference
       return logsoftmax(last);
     }
-    */;
   }
 };
 }
diff --git a/src/graph/expression_operators.cu b/src/graph/expression_operators.cu
index 0cb65c1e..be40a0d4 100644
--- a/src/graph/expression_operators.cu
+++ b/src/graph/expression_operators.cu
@@ -333,81 +333,81 @@ Expr shift(Expr a, Shape shift) {
 //  return Expression<LexicalProbNodeOp>(logits, att, eps, lf);
 //}
 
-//Expr avg_pooling(
-//    Expr x,
-//    int height,
-//    int width,
-//    int padHeight,
-//    int padWidth,
-//    int strideHeight,
-//    int strideWidth) {
-//  return Expression<PoolingOp>(x,
-//      height,
-//      width,
-//      padHeight,
-//      padWidth,
-//      strideHeight,
-//      strideWidth,
-//      "avg");
-//}
-//
-//Expr max_pooling(
-//    Expr x,
-//    int height,
-//    int width,
-//    int padHeight,
-//    int padWidth,
-//    int strideHeight,
-//    int strideWidth)
-//{
-//  return Expression<PoolingOp>(x,
-//      height,
-//      width,
-//      padHeight,
-//      padWidth,
-//      strideHeight,
-//      strideWidth,
-//      "max");
-//}
-//
-//Expr convert2cudnnFormat(Expr x) {
-//  int numWords = x->shape()[0];
-//  int numExamples = x->shape()[1];
-//  int embSize = x->shape()[2];
-//
-//  std::vector<size_t> newIndeces;
-//  for (int b = 0; b < numExamples; ++b) {
-//    for (int t = 0; t < numWords; ++t) {
-//      newIndeces.push_back((t * numExamples) + b);
-//    }
-//  }
-//
-//  auto xRows = reshape(x, {x->shape()[0] * x ->shape()[1], x->shape()[2]});
-//
-//  Shape outShape({numExamples, 1, numWords, embSize});
-//  return reshape(rows(xRows, newIndeces), outShape);
-//}
-//
-//Expr convertFromcudnnFormat(Expr x) {
-//  int batchDim = x->shape()[0];
-//  int sentenceDim = x->shape()[2];
-//  int embSize = x->shape()[3];
-//
-//  auto reshapedX = reshape(x, {batchDim * sentenceDim, embSize});
-//
-//  std::vector<size_t> newIndeces;
-//  for (int t = 0; t < sentenceDim; ++t) {
-//    for (int b = 0; b < batchDim; ++b) {
-//      newIndeces.push_back(b * sentenceDim + t);
-//    }
-//  }
-//
-//  Shape shape({batchDim, sentenceDim, embSize});
-//  return reshape(rows(reshapedX, newIndeces), shape);
-//}
-//
-//Expr pooling_with_masking(Expr x, Expr mask, int width, bool isEven) {
-//  return Expression<PoolingWithMaskingOp>(x, mask, width, isEven);
-//}
+Expr avg_pooling(
+    Expr x,
+    int height,
+    int width,
+    int padHeight,
+    int padWidth,
+    int strideHeight,
+    int strideWidth) {
+  return Expression<PoolingOp>(x,
+      height,
+      width,
+      padHeight,
+      padWidth,
+      strideHeight,
+      strideWidth,
+      "avg");
+}
+
+Expr max_pooling(
+    Expr x,
+    int height,
+    int width,
+    int padHeight,
+    int padWidth,
+    int strideHeight,
+    int strideWidth)
+{
+  return Expression<PoolingOp>(x,
+      height,
+      width,
+      padHeight,
+      padWidth,
+      strideHeight,
+      strideWidth,
+      "max");
+}
+
+Expr convert2cudnnFormat(Expr x) {
+  int numWords = x->shape()[0];
+  int numExamples = x->shape()[1];
+  int embSize = x->shape()[2];
+
+  std::vector<size_t> newIndeces;
+  for (int b = 0; b < numExamples; ++b) {
+    for (int t = 0; t < numWords; ++t) {
+      newIndeces.push_back((t * numExamples) + b);
+    }
+  }
+
+  auto xRows = reshape(x, {x->shape()[0] * x ->shape()[1], x->shape()[2]});
+
+  Shape outShape({numExamples, 1, numWords, embSize});
+  return reshape(rows(xRows, newIndeces), outShape);
+}
+
+Expr convertFromcudnnFormat(Expr x) {
+  int batchDim = x->shape()[0];
+  int sentenceDim = x->shape()[2];
+  int embSize = x->shape()[3];
+
+  auto reshapedX = reshape(x, {batchDim * sentenceDim, embSize});
+
+  std::vector<size_t> newIndeces;
+  for (int t = 0; t < sentenceDim; ++t) {
+    for (int b = 0; b < batchDim; ++b) {
+      newIndeces.push_back(b * sentenceDim + t);
+    }
+  }
+
+  Shape shape({batchDim, sentenceDim, embSize});
+  return reshape(rows(reshapedX, newIndeces), shape);
+}
+
+Expr pooling_with_masking(Expr x, Expr mask, int width, bool isEven) {
+  return Expression<PoolingWithMaskingOp>(x, mask, width, isEven);
+}
 
 }
diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h
index 9b0655ec..a2a47a61 100644
--- a/src/graph/node_operators_binary.h
+++ b/src/graph/node_operators_binary.h
@@ -5,6 +5,7 @@
 #include "graph/node.h"
 #include "functional/functional.h"
 #include "tensors/tensor_operators.h"
+#include "tensors/gpu/cudnn_wrappers.h"
 
 namespace marian {
 
@@ -743,46 +744,46 @@ struct HighwayNodeOp : public NaryNodeOp {
   const std::string type() { return "highway"; }
 };
 
-//class ConvolutionOp : public NaryNodeOp {
-//public:
-//  ConvolutionOp(
-//      const std::vector<Expr>& nodes,
-//      int hPad = 0,
-//      int wPad = 0,
-//      int hStride = 1,
-//      int wStride = 1)
-//    : NaryNodeOp(nodes),
-//      conv_(nodes[1]->shape(),
-//            nodes[2]->shape(),
-//            hPad,
-//            wPad,
-//            hStride,
-//            wStride) {
-//    conv_.getOutputShape(nodes[0]->shape(), shape_);
-//  }
-//
-//  NodeOps forwardOps() {
-//    return {NodeOp(conv_.forward(
-//          child(0)->val(),
-//          child(1)->val(),
-//          child(2)->val(),
-//          val_))};
-//  }
-//
-//  NodeOps backwardOps() {
-//    return {NodeOp(conv_.backward(
-//          child(0)->val(),
-//          child(0)->grad(),
-//          child(1)->val(),
-//          child(1)->grad(),
-//          child(2)->grad(),
-//          adj_))};
-//  }
-//
-//  const std::string type() { return "layer_convolution"; }
-//
-//protected:
-//  ConvolutionWrapper conv_;
-//};
+class ConvolutionOp : public NaryNodeOp {
+public:
+  ConvolutionOp(
+      const std::vector<Expr>& nodes,
+      int hPad = 0,
+      int wPad = 0,
+      int hStride = 1,
+      int wStride = 1)
+    : NaryNodeOp(nodes),
+      conv_(nodes[1]->shape(),
+            nodes[2]->shape(),
+            hPad,
+            wPad,
+            hStride,
+            wStride) {
+    conv_.getOutputShape(nodes[0]->shape(), shape_);
+  }
+
+  NodeOps forwardOps() {
+    return {NodeOp(conv_.forward(
+          child(0)->val(),
+          child(1)->val(),
+          child(2)->val(),
+          val_))};
+  }
+
+  NodeOps backwardOps() {
+    return {NodeOp(conv_.backward(
+          child(0)->val(),
+          child(0)->grad(),
+          child(1)->val(),
+          child(1)->grad(),
+          child(2)->grad(),
+          adj_))};
+  }
+
+  const std::string type() { return "layer_convolution"; }
+
+protected:
+  ConvolutionWrapper conv_;
+};
 
 }
diff --git a/src/graph/node_operators_unary.h b/src/graph/node_operators_unary.h
index 07c06fda..0a76471b 100644
--- a/src/graph/node_operators_unary.h
+++ b/src/graph/node_operators_unary.h
@@ -1055,84 +1055,84 @@ struct ShiftNodeOp : public UnaryNodeOp {
 //  Ptr<sparse::CSR> lf_;
 //};
 
-//class PoolingOp : public UnaryNodeOp {
-//public:
-//  PoolingOp(Expr x,
-//            int height,
-//            int width,
-//            int padHeight,
-//            int padWidth,
-//            int strideHeight,
-//            int strideWidth,
-//            std::string mode)
-//      : UnaryNodeOp(x),
-//        pooling_(height,
-//                 width,
-//                 padHeight,
-//                 padWidth,
-//                 strideHeight,
-//                 strideWidth,
-//                 mode) {
-//  }
-//
-//  NodeOps forwardOps() {
-//    return {NodeOp(pooling_.forward(child(0)->val(), val_))};
-//  }
-//
-//  NodeOps backwardOps() {
-//    return {NodeOp(pooling_.backward(
-//          child(0)->val(),
-//          child(0)->grad(),
-//          val_,
-//          adj_))};
-//  }
-//
-//  const std::string type() { return "layer_pooling"; }
-//
-//
-//protected:
-//  PoolingWrapper pooling_;
-//};
-//
-//class PoolingWithMaskingOp : public UnaryNodeOp {
-//  public:
-//    PoolingWithMaskingOp( Expr x, Expr mask, int width, bool isEven=false)
-//      : UnaryNodeOp(x),
-//        mask_(mask),
-//        width_(width),
-//        isEven_(isEven)
-//    {
-//      auto xShape = x->shape();
-//      int dimBatch = xShape[0];
-//      int dimWord = xShape[1];
-//      int cols = (isEven_) ? xShape[2] - 1 : xShape[2];
-//      int dimSentence = (cols / width_) + (cols % width_ != 0);
-//      shape_ = {dimBatch, dimWord, dimSentence};
-//    }
-//
-//    NodeOps forwardOps() {
-//      return {NodeOp(PoolingWithMaskingForward(val_,
-//                                              child(0)->val(),
-//                                              mask_->val(),
-//                                              width_,
-//                                              isEven_))};
-//    }
-//
-//    NodeOps backwardOps() {
-//      return {NodeOp(PoolingWithMaskingBackward(adj_,
-//                                               child(0)->grad(),
-//                                               child(0)->val(),
-//                                               mask_->val(),
-//                                               width_,
-//                                               isEven_))};
-//    }
-//
-//    const std::string type() {return "layer_pooling";}
-//
-//  protected:
-//    Expr mask_;
-//    int width_;
-//    bool isEven_;
-//};
+class PoolingOp : public UnaryNodeOp {
+public:
+  PoolingOp(Expr x,
+            int height,
+            int width,
+            int padHeight,
+            int padWidth,
+            int strideHeight,
+            int strideWidth,
+            std::string mode)
+      : UnaryNodeOp(x),
+        pooling_(height,
+                 width,
+                 padHeight,
+                 padWidth,
+                 strideHeight,
+                 strideWidth,
+                 mode) {
+  }
+
+  NodeOps forwardOps() {
+    return {NodeOp(pooling_.forward(child(0)->val(), val_))};
+  }
+
+  NodeOps backwardOps() {
+    return {NodeOp(pooling_.backward(
+          child(0)->val(),
+          child(0)->grad(),
+          val_,
+          adj_))};
+  }
+
+  const std::string type() { return "layer_pooling"; }
+
+
+protected:
+  PoolingWrapper pooling_;
+};
+
+class PoolingWithMaskingOp : public UnaryNodeOp {
+  public:
+    PoolingWithMaskingOp( Expr x, Expr mask, int width, bool isEven=false)
+      : UnaryNodeOp(x),
+        mask_(mask),
+        width_(width),
+        isEven_(isEven)
+    {
+      auto xShape = x->shape();
+      int dimBatch = xShape[0];
+      int dimWord = xShape[1];
+      int cols = (isEven_) ? xShape[2] - 1 : xShape[2];
+      int dimSentence = (cols / width_) + (cols % width_ != 0);
+      shape_ = {dimBatch, dimWord, dimSentence};
+    }
+
+    NodeOps forwardOps() {
+      return {NodeOp(PoolingWithMaskingForward(val_,
+                                              child(0)->val(),
+                                              mask_->val(),
+                                              width_,
+                                              isEven_))};
+    }
+
+    NodeOps backwardOps() {
+      return {NodeOp(PoolingWithMaskingBackward(adj_,
+                                               child(0)->grad(),
+                                               child(0)->val(),
+                                               mask_->val(),
+                                               width_,
+                                               isEven_))};
+    }
+
+    const std::string type() {return "layer_pooling";}
+
+  protected:
+    Expr mask_;
+    int width_;
+    bool isEven_;
+};
 
 }
diff --git a/src/layers/convolution.cu b/src/layers/convolution.cu
index 83e881bf..b0749450 100644
--- a/src/layers/convolution.cu
+++ b/src/layers/convolution.cu
@@ -25,12 +25,11 @@ Expr Convolution::apply(Expr x) {
                             keywords::init=inits::zeros);
 
   std::vector<Expr> nodes = {x, kernel, bias};
-  ABORT("Temporarily not implemented");
-  //return Expression<ConvolutionOp>(nodes,
-  //                                 paddings.first,
-  //                                 paddings.second,
-  //                                 strides.first,
-  //                                 strides.second);
+  return Expression<ConvolutionOp>(nodes,
+                                   paddings.first,
+                                   paddings.second,
+                                   strides.first,
+                                   strides.second);
 }
 
 Expr Convolution::apply(const std::vector<Expr>&) {
diff --git a/src/models/model_factory.cpp b/src/models/model_factory.cpp
index 0fe58c3b..ddeb1071 100644
--- a/src/models/model_factory.cpp
+++ b/src/models/model_factory.cpp
@@ -8,13 +8,13 @@
 #include "models/nematus.h"
 #include "models/encdec.h"
 
-#ifdef USE_CUDNN
+#ifdef CUDNN
 #include "models/char_s2s.h"
 #endif
 
 #ifdef COMPILE_EXAMPLES
 #include "examples/mnist/model.h"
-#ifdef USE_CUDNN
+#ifdef CUDNN
 #include "examples/mnist/model_lenet.h"
 #endif
 #endif
@@ -26,7 +26,7 @@ Ptr<EncoderBase> EncoderFactory::construct() {
   if(options_->get<std::string>("type") == "s2s")
     return New<EncoderS2S>(options_);
 
-#ifdef USE_CUDNN
+#ifdef CUDNN
   if(options_->get<std::string>("type") == "char-s2s")
     return New<CharS2SEncoder>(options_);
 #endif
diff --git a/src/tensors/tensor_operators.h b/src/tensors/tensor_operators.h
index 06f76b64..30c42c4a 100644
--- a/src/tensors/tensor_operators.h
+++ b/src/tensors/tensor_operators.h
@@ -209,5 +209,8 @@ namespace marian {
       return cpu::L2Norm(in);
     }
   }
-
+  
+  DISPATCH5(PoolingWithMaskingForward, marian::Tensor, marian::Tensor, marian::Tensor, int, bool)
+  DISPATCH6(PoolingWithMaskingBackward, marian::Tensor, marian::Tensor, marian::Tensor, marian::Tensor, int, bool)
+  
 }