Enable 'empty' function for mkldnn

PenghuiCheng · facebook-github-bot · commit 57f932a6384e · 2019-06-04T14:16:13.000-07:00
Summary: Pull Request resolved: #21184 Differential Revision: D15625296 Pulled By: bddppq fbshipit-source-id: 47d26798bcf48e227ffd813f299959a7b8993641
diff --git a/aten/src/ATen/native/mkldnn/MKLDNNCommon.cpp b/aten/src/ATen/native/mkldnn/MKLDNNCommon.cpp
@@ -50,15 +50,6 @@ Tensor new_with_itensor_mkldnn(ideep::tensor&& it, const TensorOptions& options)
     std::vector<int64_t>(dims.begin(), dims.end()));
 }
 
-Tensor new_with_sizes_mkldnn(IntArrayRef sizes, 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
-  ideep::tensor::dims dst_dims (sizes.begin(), sizes.end());
-  ideep::tensor it;
-  it.resize<AllocForMKLDNN>(dst_dims, ideep::tensor::data_type::f32);
-  return new_with_itensor_mkldnn(std::move(it), options);
-}
-
 ideep::tensor& itensor_from_mkldnn(const MKLDNNTensor& mkldnn_tensor) {
   AT_ASSERTM(mkldnn_tensor.is_mkldnn(),
              "mkldnn_to_dense expects MKL-DNN tensor input");
diff --git a/aten/src/ATen/native/mkldnn/MKLDNNCommon.h b/aten/src/ATen/native/mkldnn/MKLDNNCommon.h
@@ -26,9 +26,6 @@ struct AllocForMKLDNN {
 // Construct aten MKL-DNN tensor given an ideep tensor
 Tensor new_with_itensor_mkldnn(ideep::tensor&& it, const TensorOptions& options);
 
-// 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
 ideep::tensor& itensor_from_mkldnn(const Tensor& mkldnn_tensor);
 
diff --git a/aten/src/ATen/native/mkldnn/MKLDNNConversions.cpp b/aten/src/ATen/native/mkldnn/MKLDNNConversions.cpp
@@ -24,9 +24,11 @@ Tensor dense_to_mkldnn(const Tensor& cpu_tensor) {
              "dense_to_mkldnn expects dense CPU tensor input");
   AT_ASSERTM(cpu_tensor.scalar_type() == ScalarType::Float,
              "dense_to_mkldnn expects float tensor input");
+  AT_ASSERTM(cpu_tensor.dim() <= 5,
+             "Can't convert cpu tensor with the number of dimensions > 5");
   // TODO: consider to convert non-contiguous tensor to `ideep::tensor` directly.
   auto cpu_tensor_cont = cpu_tensor.contiguous();
-  Tensor mkldnn_tensor = new_with_sizes_mkldnn(cpu_tensor_cont.sizes(), cpu_tensor_cont.options());
+  Tensor mkldnn_tensor = empty_mkldnn(cpu_tensor_cont.sizes(), cpu_tensor_cont.options());
   ideep::tensor& dtensor = itensor_from_mkldnn(mkldnn_tensor);
   dtensor.reorder_from(dtensor.get_dims(),
                        ideep::tensor::data_type::f32,
diff --git a/aten/src/ATen/native/mkldnn/TensorFactories.cpp b/aten/src/ATen/native/mkldnn/TensorFactories.cpp
@@ -0,0 +1,24 @@
+#include <ATen/native/mkldnn/MKLDNNCommon.h>
+
+namespace at { namespace native {
+
+#if AT_MKLDNN_ENABLED()
+
+Tensor empty_mkldnn(IntArrayRef sizes, 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
+  ideep::tensor::dims dst_dims (sizes.begin(), sizes.end());
+  ideep::tensor it;
+  it.resize<AllocForMKLDNN>(dst_dims, ideep::tensor::data_type::f32);
+  return new_with_itensor_mkldnn(std::move(it), options);
+}
+
+#else
+
+Tensor empty_mkldnn(IntArrayRef sizes, const TensorOptions& options) {
+  AT_ERROR("empty_mkldnn: MKL-DNN build is disabled");
+}
+
+#endif // AT_MKLDNN_ENABLED()
+
+}}
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -688,6 +688,7 @@
   dispatch:
     CPU: empty_cpu
     CUDA: empty_cuda
+    MkldnnCPU: empty_mkldnn
     SparseCPU: empty_sparse
     SparseCUDA: empty_sparse
 
diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -45,7 +45,7 @@ def test_unsupported(self):
             with self.assertRaises(RuntimeError) as context:
                 torch.randn(1, 2, 3, 4, dtype=torch.float, device=torch.device('cuda')).to_mkldnn()
         # some factory functions
-        for creator in [torch.empty, torch.ones, torch.zeros, torch.randn, torch.rand]:
+        for creator in [torch.ones, torch.zeros, torch.randn, torch.rand]:
             with self.assertRaises(RuntimeError) as context:
                 creator(1, 2, 3, 4, dtype=torch.float, device=torch.device('cpu'), layout=torch._mkldnn)
 
@@ -289,6 +289,11 @@ def test_set_data_tensorimpl_type(self):
         with self.assertRaisesRegex(RuntimeError, 'different types of TensorImpl'):
             x.data = x_mkldnn
 
+    def test_empty(self):
+        x1 = torch.empty(4, 5, 2, 3, dtype=torch.float32)
+        x2 = torch.empty(4, 5, 2, 3, dtype=torch.float32, layout=torch._mkldnn)
+        self.assertEqual(x1.size(), x2.to_dense().size())
+        self.assertEqual(x1.dtype, x2.to_dense().dtype)
 
 if __name__ == '__main__':
     run_tests()
