pytorch · bddppq · Apr 12, 2019 · Apr 13, 2019 · Apr 13, 2019 · Apr 16, 2019
diff --git a/aten/src/ATen/native/LegacyDefinitions.cpp b/aten/src/ATen/native/LegacyDefinitions.cpp
@@ -2,6 +2,8 @@
 #include <ATen/NativeFunctions.h>
 #include <ATen/LegacyTHFunctions.h>
 
+#include <ATen/native/mkldnn/TensorShape.h>
+
 namespace at { namespace native {
 
 // Methods
@@ -63,6 +65,9 @@ Tensor & masked_scatter_(Tensor& self, const Tensor & mask, const Tensor & sourc
 }
 
 Tensor view(const Tensor& self, IntArrayRef size) {
+  if (self.is_mkldnn()) {
+    return mkldnn_view(self, size);
+  }
   return at::legacy::th::_th_view(self, size);
 }
 

diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -431,6 +431,11 @@ Tensor reshape(const Tensor& self, IntArrayRef proposed_shape) {
     AT_ERROR("reshape is not implemented for sparse tensors");
   }
   auto shape = infer_size(proposed_shape, self.numel());
+
+  if (self.is_mkldnn()) {
+    return at::mkldnn_reshape(self, shape);
+  }
+
   if (auto stride = THTensor_compute_stride(self.sizes(), self.strides(), shape)) {
     return self.as_strided(shape, *stride);
   }

diff --git a/aten/src/ATen/native/mkldnn/TensorShape.cpp b/aten/src/ATen/native/mkldnn/TensorShape.cpp
@@ -0,0 +1,46 @@
+#include <ATen/ATen.h>
+#include <ATen/Config.h>
+#include <ATen/InferSize.h>
+#include <ATen/NativeFunctions.h>
+
+#if !AT_MKLDNN_ENABLED()
+
+namespace at {
+namespace native {
+
+Tensor mkldnn_view(const Tensor& self, IntArrayRef size) {
+  AT_ERROR("mkldnn_reshape: ATen not compiled with MKLDNN support");
+}
+
+Tensor mkldnn_reshape(const Tensor& self, IntArrayRef size) {
+  AT_ERROR("mkldnn_reshape: ATen not compiled with MKLDNN support");
+}
+
+} // namespace native
+} // namespace at
+
+#else // AT_MKLDNN_EBABLED
+
+#include <ATen/native/mkldnn/MKLDNNCommon.h>
+
+namespace at {
+namespace native {
+
+Tensor mkldnn_view(const Tensor& self, IntArrayRef size) {
+  AT_ERROR(
+      "Currently Mkldnn tensor does not support view. Change to use reshape instead");
+}
+
+Tensor mkldnn_reshape(const Tensor& self, IntArrayRef size) {
+  auto inferred_size = at::infer_size(size, self.numel());
+  const ideep::tensor& x = itensor_from_mkldnn(self);
+  ideep::tensor y;
+  ideep::direct_copy::compute<AllocForMKLDNN>(x, y);
+  y.reshape({inferred_size.cbegin(), inferred_size.cend()});
+  return new_with_itensor_mkldnn(std::move(y), self.options());
+}
+
+} // namespace native
+} // namespace at
+
+#endif // AT_MKLDNN_EBABLED
diff --git a/aten/src/ATen/native/mkldnn/TensorShape.h b/aten/src/ATen/native/mkldnn/TensorShape.h
@@ -0,0 +1,11 @@
+#pragma once
+
+#include <ATen/ATen.h>
+
+namespace at {
+namespace native {
+
+Tensor mkldnn_view(const Tensor& self, IntArrayRef size);
+
+} // namespace native
+} // namespace at
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1522,6 +1522,12 @@
   variants: function, method
   device_guard: False
 
+- func: mkldnn_reshape(Tensor self, int[] shape) -> Tensor
+  device_guard: False
+  requires_tensor: True
+  dispatch:
+    MkldnnCPU: mkldnn_reshape
+
 - func: reshape_as(Tensor self, Tensor other) -> Tensor
   variants: method
   device_guard: False

diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -192,6 +192,21 @@ def test_add(self):
         torch.add(mx, my, alpha=alpha, out=mkldnn_out)
         self.assertEqual(out, mkldnn_out.to_dense())
 
+    def test_view(self):
+        x = torch.randn(3, 4, 5, dtype=torch.float32).to_mkldnn()
+        self.assertRaisesRegex(RuntimeError,
+                               "Change to use reshape",
+                               lambda: x.view(x.size(0), -1))
+
+    def test_reshape(self):
+        x = torch.randn(3, 4, 5, dtype=torch.float32) * 10
+        size = (x.size(0), -1)
+
+        self.assertEqual(
+            x.reshape(size),
+            x.to_mkldnn().reshape(size).to_dense(),
+        )
+
 
 if __name__ == '__main__':
     run_tests()