pytorch · bddppq · Apr 12, 2019 · Apr 13, 2019 · Apr 13, 2019 · Apr 13, 2019
diff --git a/aten/src/ATen/native/Pooling.cpp b/aten/src/ATen/native/Pooling.cpp
@@ -114,6 +114,10 @@ Tensor max_pool2d(
     IntArrayRef padding,
     IntArrayRef dilation,
     bool ceil_mode) {
+  if (self.is_mkldnn()) {
+    return at::mkldnn_max_pool2d(
+        self, kernel_size, stride, padding, dilation, ceil_mode);
+  }
   auto output_and_indices = at::max_pool2d_with_indices(
       self, kernel_size, stride, padding, dilation, ceil_mode);
   return std::get<0>(output_and_indices);

diff --git a/aten/src/ATen/native/mkldnn/MKLDNNCommon.cpp b/aten/src/ATen/native/mkldnn/MKLDNNCommon.cpp
@@ -43,7 +43,6 @@ using MKLDNNTensor = Tensor;
 Tensor new_with_itensor_mkldnn(ideep::tensor&& it, const TensorOptions& options) {
   // NOTE: int32_t dims from ideep::tensor but sizes needs int64_t
   // TODO: support int64_t dims in ideep::tensor to avoid extra conversion
-  AT_ASSERT(!it.has_extra());
   auto dims = it.get_dims();
   IDeepTensorWrapperPtr handle = c10::make_intrusive<IDeepTensorWrapper>(std::move(it));
   return detail::make_tensor<MKLDNNTensorImpl>(

diff --git a/aten/src/ATen/native/mkldnn/MKLDNNCommon.h b/aten/src/ATen/native/mkldnn/MKLDNNCommon.h
@@ -23,10 +23,10 @@ struct AllocForMKLDNN {
   }
 };
 
-// Construct MKL-DNN tensor given an ideep tensor
+// Construct aten MKL-DNN tensor given an ideep tensor
 Tensor new_with_itensor_mkldnn(ideep::tensor&& it, const TensorOptions& options);
 
-// Construct MKL-DNN tensor given `sizes` for allocation
+// Construct aten MKL-DNN tensor given `sizes` for allocation
 Tensor new_with_sizes_mkldnn(IntArrayRef sizes, const TensorOptions& options);
 
 // Retrieve `ideep::tensor` from MKL-DNN tensor

diff --git a/aten/src/ATen/native/mkldnn/Pooling.cpp b/aten/src/ATen/native/mkldnn/Pooling.cpp
@@ -0,0 +1,142 @@
+#include <ATen/ATen.h>
+#include <ATen/Config.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/native/utils/ParamUtils.h>
+#include <tuple>
+
+
+#if !AT_MKLDNN_ENABLED()
+
+namespace at {
+namespace native {
+
+Tensor mkldnn_max_pool2d(
+    const Tensor& self,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode) {
+  AT_ERROR(
+      "mkldnn_max_pool2d: ATen not compiled with MKLDNN support");
+}
+
+Tensor mkldnn_avg_pool2d(
+    const Tensor& self,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad) {
+  AT_ERROR("mkldnn_avg_pool2d: ATen not compiled with MKLDNN support");
+}
+
+Tensor& mkldnn_avg_pool2d_out(
+    Tensor& output,
+    const Tensor& self,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad) {
+  AT_ERROR("mkldnn_avg_pool2d_out: ATen not compiled with MKLDNN support");
+}
+} // namespace native
+} // namespace at
+
+#else // AT_MKLDNN_ENABLED
+
+#include <ATen/native/mkldnn/MKLDNNCommon.h>
+#include <ATen/native/mkldnn/Utils.h>
+
+namespace at {
+namespace native {
+
+static Tensor _mkldnn_pool2d(
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode,
+    ideep::algorithm algo) {
+  AT_CHECK(!ceil_mode, "Currently Mkldnn Pooling operators do not support ceil_mode.");
+  auto kernel_size_vec = expand_param_if_needed(kernel_size, "kernel_size", 2);
+  auto stride_vec = expand_param_if_needed(stride, "stride", 2);
+  auto padding_vec = expand_param_if_needed(padding, "padding", 2);
+  auto dilation_vec = expand_param_if_needed(dilation, "dilation", 2);
+
+  const ideep::tensor& x = itensor_from_mkldnn(input);
+  const std::vector<int64_t> output_sizes = pool_output_sizes(
+      input.sizes(),
+      kernel_size_vec,
+      stride_vec,
+      padding_vec,
+      dilation_vec,
+      ceil_mode);
+  ideep::tensor y;
+  ideep::pooling_forward::compute<AllocForMKLDNN>(
+      x,
+      {output_sizes.cbegin(), output_sizes.cend()},
+      y,
+      {stride_vec.cbegin(), stride_vec.cend()},
+      {kernel_size_vec.cbegin(), kernel_size_vec.cend()},
+      {padding_vec.cbegin(), padding_vec.cend()},
+      {padding_vec.cbegin(), padding_vec.cend()},
+      algo,
+      ideep::prop_kind::forward);
+
+  return new_with_itensor_mkldnn(std::move(y), input.options());
+}
+
+Tensor mkldnn_max_pool2d(
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode) {
+  return _mkldnn_pool2d(
+      input,
+      kernel_size,
+      stride,
+      padding,
+      dilation,
+      ceil_mode,
+      ideep::algorithm::pooling_max);
+}
+
+Tensor mkldnn_avg_pool2d(
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad) {
+  return _mkldnn_pool2d(
+      input,
+      kernel_size,
+      stride,
+      padding,
+      std::vector<int64_t>{1, 1},
+      ceil_mode,
+      count_include_pad ? ideep::algorithm::pooling_avg_include_padding
+                        : ideep::algorithm::pooling_avg_exclude_padding);
+}
+
+Tensor& mkldnn_avg_pool2d_out(
+    Tensor& output,
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad) {
+  AT_ERROR(
+      "mkldnn_avg_pool2d_out: in-place mkldnn operations are not supported yet");
+}
+
+} // namespace native
+} // namespace at
+
+#endif // AT_MKLDNN_ENABLED
diff --git a/aten/src/ATen/native/mkldnn/Relu.cpp b/aten/src/ATen/native/mkldnn/Relu.cpp
@@ -0,0 +1,43 @@
+#include <ATen/ATen.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/Config.h>
+
+
+#if !AT_MKLDNN_ENABLED()
+
+namespace at { namespace native {
+
+Tensor mkldnn_relu(const Tensor& input) {
+  AT_ERROR("mkldnn_relu: ATen not compiled with MKLDNN support");
+}
+
+Tensor& mkldnn_relu_(Tensor& input) {
+  AT_ERROR("mkldnn_relu_: ATen not compiled with MKLDNN support");
+}
+
+}}
+
+#else // AT_MKLDNN_EBABLED
+
+#include <ATen/native/mkldnn/MKLDNNCommon.h>
+
+namespace at { namespace native {
+
+Tensor mkldnn_relu(const Tensor& input) {
+  const ideep::tensor& x = itensor_from_mkldnn(input);
+  ideep::tensor y;
+  ideep::eltwise_forward::compute<AllocForMKLDNN>(
+      x, y, ideep::algorithm::eltwise_relu, ideep::prop_kind::forward_training, /*alpha*/ 0.0);
+  return new_with_itensor_mkldnn(std::move(y), input.options());
+}
+
+Tensor& mkldnn_relu_(Tensor& input) {
+  ideep::tensor& x = itensor_from_mkldnn(input);
+  ideep::eltwise_forward::compute<AllocForMKLDNN>(
+      x, x, ideep::algorithm::eltwise_relu, ideep::prop_kind::forward_training, /*alpha*/ 0.0);
+  return input;
+}
+
+}}
+
+#endif // AT_MKLDNN_EBABLED
diff --git a/aten/src/ATen/native/mkldnn/Utils.cpp b/aten/src/ATen/native/mkldnn/Utils.cpp
@@ -1,4 +1,5 @@
 #include <ATen/native/mkldnn/Utils.h>
+#include <THNN/generic/pooling_shape.h>
 
 namespace at { namespace native {
 
@@ -20,4 +21,30 @@ std::vector<int64_t> conv_output_size(
   return output_size;
 }
 
+std::vector<int64_t> pool_output_sizes(
+    IntArrayRef input_size,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode) {
+  std::vector<int64_t> output_size(input_size.size());
+  // copy N and C
+  output_size[0] = input_size[0];
+  output_size[1] = input_size[1];
+
+  for (int i = 2; i < input_size.size(); ++i) {
+    output_size[i] = pooling_output_shape<int64_t>(
+      input_size[i],
+      kernel_size[i - 2],
+      padding[i - 2],
+      stride[i - 2],
+      dilation[i - 2],
+      ceil_mode
+    );
+  }
+
+   return output_size;
+}
+
 }}
diff --git a/aten/src/ATen/native/mkldnn/Utils.h b/aten/src/ATen/native/mkldnn/Utils.h
@@ -11,4 +11,12 @@ std::vector<int64_t> conv_output_size(
     IntArrayRef padding,
     IntArrayRef stride,
     IntArrayRef dilation);
+
+std::vector<int64_t> pool_output_sizes(
+    IntArrayRef input_size,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode);
 }}
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1141,6 +1141,11 @@
 
 - func: max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor
 
+- func: mkldnn_max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor
+  requires_tensor: True
+  dispatch:
+    MkldnnCPU: mkldnn_max_pool2d
+
 - func: max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor
 
 # FIXME: These could be combined as optional<ScalarType> but for https://github.com/pytorch/pytorch/issues/6593.
@@ -1522,9 +1527,17 @@
 
 - func: relu(Tensor self) -> Tensor
   variants: function, method
+  dispatch:
+    CPU: relu
+    CUDA: relu
+    MkldnnCPU: mkldnn_relu
 
 - func: relu_(Tensor(a!) self) -> Tensor(a!)
   variants: function, method
+  dispatch:
+    CPU: relu_
+    CUDA: relu_
+    MkldnnCPU: mkldnn_relu_
 
 - func: prelu(Tensor self, Tensor weight) -> Tensor
   variants: function, method
@@ -3670,9 +3683,17 @@
 
 - func: avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, *, Tensor(a!) out) -> Tensor(a!)
   python_module: nn
+  dispatch:
+    CPU: avg_pool2d_out
+    CUDA: avg_pool2d_out
+    MkldnnCPU: mkldnn_avg_pool2d_out
 
 - func: avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True) -> Tensor
   python_module: nn
+  dispatch:
+    CPU: avg_pool2d
+    CUDA: avg_pool2d
+    MkldnnCPU: mkldnn_avg_pool2d
 
 - func: avg_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, bool ceil_mode, bool count_include_pad, *, Tensor(a!) grad_input) -> Tensor(a!)
   python_module: nn

diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -109,6 +109,45 @@ def test_conv2d(self):
                     conv2d(x),
                     mkldnn_conv2d(x.to_mkldnn()).to_dense())
 
+    def test_relu(self):
+        x = torch.randn((4, 5), dtype=torch.float32) * 10
+        self.assertEqual(torch.relu(x), torch.relu(x.to_mkldnn()).to_dense())
+
+    def test_relu_(self):
+        x1 = torch.randn((4, 5), dtype=torch.float32) * 10
+        x2 = x1.clone().to_mkldnn()
+        self.assertEqual(torch.relu_(x1), torch.relu_(x2).to_dense())
+
+    def test_max_pool2d(self):
+        N = torch.randint(3, 10, (1,)).item()
+        C = torch.randint(3, 10, (1,)).item()
+        x = torch.randn(N, C, 64, 64, dtype=torch.float32) * 10
+
+        max_pool2d = torch.nn.MaxPool2d(
+            kernel_size=3,
+            stride=2,
+            padding=1)
+
+        self.assertEqual(
+            max_pool2d(x),
+            max_pool2d(x.to_mkldnn()).to_dense())
+
+    def test_avg_pool2d(self):
+        N = torch.randint(3, 10, (1,)).item()
+        C = torch.randint(3, 10, (1,)).item()
+        x = torch.randn(N, C, 64, 64, dtype=torch.float32) * 10
+
+        for count_include_pad in [True, False]:
+            avg_pool2d = torch.nn.AvgPool2d(
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                count_include_pad=count_include_pad)
+
+            self.assertEqual(
+                avg_pool2d(x),
+                avg_pool2d(x.to_mkldnn()).to_dense())
+
 
 if __name__ == '__main__':
     run_tests()
diff --git a/torch/nn/modules/activation.py b/torch/nn/modules/activation.py
@@ -60,7 +60,7 @@ def extra_repr(self):
 
 
 @weak_module
-class ReLU(Threshold):
+class ReLU(Module):
     r"""Applies the rectified linear unit function element-wise:
 
     :math:`\text{ReLU}(x)= \max(0, x)`
@@ -88,9 +88,15 @@ class ReLU(Threshold):
         >>> input = torch.randn(2).unsqueeze(0)
         >>> output = torch.cat((m(input),m(-input)))
     """
+    __constants__ = ['inplace']
 
     def __init__(self, inplace=False):
-        super(ReLU, self).__init__(0., 0., inplace)
+        super(ReLU, self).__init__()
+        self.inplace = inplace
+
+    @weak_script_method
+    def forward(self, input):
+        return F.relu(input, inplace=self.inplace)
 
     def extra_repr(self):
         inplace_str = 'inplace' if self.inplace else ''