Add memory format support to the resize_ op. (#28292)

VitalyFedyunin · facebook-github-bot · commit cb43170dcb93 · 2019-11-18T05:35:44.000-08:00
Summary: Pull Request resolved: #28292 Allows to simplify patterns like: 1. output.resize_({sizeB, sizeC, osizeH, osizeW}).as_strided_({sizeB, sizeC, osizeH, osizeW}, {sizeC*osizeH*osizeW, 1, osizeW*sizeC, sizeC}); 2. output.resize_({nbatch, nInputPlane, outputHeight, outputWidth}); indices.resize_({nbatch, nInputPlane, outputHeight, outputWidth}); output.unsafeGetTensorImpl()->empty_tensor_restride(memory_format); indices.unsafeGetTensorImpl()->empty_tensor_restride(memory_format); 3. gradInput.resize_as_(input); gradInput.unsafeGetTensorImpl()->empty_tensor_restride(memory_format); Test Plan: Imported from OSS Differential Revision: D18044978 Pulled By: VitalyFedyunin fbshipit-source-id: bbf67c25f9cf88bc6e949089a3b247df50f86dc4
diff --git a/aten/src/ATen/native/Resize.cpp b/aten/src/ATen/native/Resize.cpp
@@ -5,15 +5,27 @@
 
 namespace at { namespace native {
 
-Tensor& resize_cpu_(Tensor& self, IntArrayRef size) {
+Tensor& resize_cpu_(
+    Tensor& self,
+    IntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
 #ifdef BUILD_NAMEDTENSOR
   if (self.has_names()) {
-    return resize_named_tensor_(self, size);
+    return resize_named_tensor_(self, size, optional_memory_format);
   }
 #endif
   auto* self_ = self.unsafeGetTensorImpl();
   resize_impl_cpu_(self_, size, /*strides=*/c10::nullopt);
   self_->maybe_zero_dim(size.size() == 0);
+  if (optional_memory_format.has_value()) {
+    auto memory_format =
+        optional_memory_format.value();
+    TORCH_CHECK(
+        memory_format != MemoryFormat::Preserve,
+        "Unsupported memory format",
+        memory_format);
+    self_->empty_tensor_restride(memory_format);
+  }
   return self;
 }
 
diff --git a/aten/src/ATen/native/ResizeCommon.h b/aten/src/ATen/native/ResizeCommon.h
@@ -6,16 +6,27 @@
 namespace at { namespace native {
 
 #ifdef BUILD_NAMEDTENSOR
-inline Tensor& resize_named_tensor_(Tensor& self, IntArrayRef size) {
+inline Tensor& resize_named_tensor_(
+    Tensor& self,
+    IntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
   TORCH_INTERNAL_ASSERT(self.has_names());
   TORCH_CHECK(
       self.sizes() == size,
       "Cannot resize named tensor with resize_ or resize_as_ (tried to resize "
-      "Tensor", self.names(), " with size ", self.sizes(), " to ", size,
+      "Tensor",
+      self.names(),
+      " with size ",
+      self.sizes(),
+      " to ",
+      size,
       "). This may be caused by passing a named tensor ",
       "as an `out=` argument; please ensure that the sizes are the same. ");
+  TORCH_CHECK(
+      !optional_memory_format.has_value(),
+      "Unsupported memory format for named tensor resize ",
+      optional_memory_format.value());
   return self;
 }
 #endif
-
 }}
diff --git a/aten/src/ATen/native/cuda/Resize.cu b/aten/src/ATen/native/cuda/Resize.cu
@@ -4,18 +4,31 @@
 #include <ATen/native/cuda/Resize.cuh>
 #include <ATen/native/ResizeCommon.h>
 
-namespace at { namespace native {
+namespace at {
+namespace native {
 
-Tensor& resize_cuda_(Tensor& self, IntArrayRef size) {
+Tensor& resize_cuda_(
+    Tensor& self,
+    IntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
 #ifdef BUILD_NAMEDTENSOR
   if (self.has_names()) {
-    return resize_named_tensor_(self, size);
+    return resize_named_tensor_(self, size, optional_memory_format);
   }
 #endif
   auto* self_ = self.unsafeGetTensorImpl();
   resize_impl_cuda_(self_, size, /*strides=*/c10::nullopt);
   self_->maybe_zero_dim(size.size() == 0);
+  if (optional_memory_format.has_value()) {
+    auto memory_format =
+        optional_memory_format.value();
+    TORCH_CHECK(
+        memory_format != MemoryFormat::Preserve,
+        "Unsupported memory format",
+        memory_format);
+    self_->empty_tensor_restride(memory_format);
+  }
   return self;
 }
-
-}}
+} // namespace native
+} // namespace at
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1015,7 +1015,7 @@
     CPU: empty_per_channel_affine_quantized_other_backends_stub
     QuantizedCPU: empty_per_channel_affine_quantized_cpu
 
-- func: resize_(Tensor(a!) self, int[] size) -> Tensor(a!)
+- func: resize_(Tensor(a!) self, int[] size, *, MemoryFormat? memory_format=None) -> Tensor(a!)
   supports_named_tensor: True
   variants: method
   device_guard: False
diff --git a/aten/src/ATen/native/quantized/cpu/tensor_operators.cpp b/aten/src/ATen/native/quantized/cpu/tensor_operators.cpp
@@ -59,7 +59,14 @@ AT_FORALL_OPERATORS(DEFINE_COMPARATOR)
 #undef AT_FORALL_OPERATORS
 #undef DEFINE_COMPARATOR
 
-Tensor& quantized_resize_cpu_(Tensor& self, IntArrayRef size) {
+Tensor& quantized_resize_cpu_(
+    Tensor& self,
+    IntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
+  TORCH_CHECK(
+      !optional_memory_format.has_value(),
+      "Unsupported memory format for quantized tensor resize ",
+      optional_memory_format.value());
   auto qscheme = self.quantizer()->qscheme();
   TORCH_CHECK(
       qscheme == QScheme::PER_TENSOR_AFFINE ||
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -12065,6 +12065,11 @@ def test_memory_format_resize_as(self, device):
         flat.resize_as_(nhwc, memory_format=torch.preserve_format)
         self.assertTrue(flat.is_contiguous(memory_format=torch.channels_last))
 
+    def test_memory_format_resize_(self, device):
+        flat = torch.randn(10 * 3 * 32 * 32, device=device)
+        flat.resize_((10, 3, 32, 32), memory_format=torch.channels_last)
+        self.assertTrue(flat.is_contiguous(memory_format=torch.channels_last))
+
     def test_memory_format_empty_like(self, device):
         x = torch.randn(4, 3, 8, 8, device=device)
         nhwc = x.contiguous(memory_format=torch.channels_last)
diff --git a/tools/autograd/templates/VariableType.h b/tools/autograd/templates/VariableType.h
@@ -69,7 +69,7 @@ namespace VariableType {
   int64_t output_nr(const Tensor & self);
   int64_t _version(const Tensor & self);
   Tensor & copy_(Tensor & self, const Tensor & src, bool non_blocking);
-  Tensor & resize_(Tensor & self, IntArrayRef size);
+  Tensor & resize_(Tensor & self, IntArrayRef size, c10::optional<MemoryFormat> optional_memory_format);
   Tensor & resize_as_(Tensor & self, const Tensor & the_template, c10::optional<MemoryFormat> optional_memory_format);
   Tensor detach(const Tensor & self);
   Tensor & detach_(Tensor & self);
diff --git a/torch/csrc/autograd/VariableTypeManual.cpp b/torch/csrc/autograd/VariableTypeManual.cpp
@@ -161,7 +161,10 @@ Tensor & copy_(Tensor & self, const Tensor & src, bool non_blocking) {
   return self;
 }
 
-Tensor & resize_(Tensor & self, IntArrayRef size) {
+Tensor& resize_(
+    Tensor& self,
+    IntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
   auto& self_ = unpack(self, "self", 0);
   if (as_variable_ref(self).requires_grad()) {
     AT_ERROR("cannot resize variables that require grad");
@@ -173,7 +176,7 @@ Tensor & resize_(Tensor & self, IntArrayRef size) {
   }
   {
     at::AutoNonVariableTypeMode non_var_type_mode(true);
-    self_.resize_(size);
+    self_.resize_(size, std::move(optional_memory_format));
   }
   return self;
 }

Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,10 @@ Tensor & copy_(Tensor & self, const Tensor & src, bool non_blocking) {`
`161`	`161`	`return self;`
`162`	`162`	`}`
`163`	`163`
`164`		`-Tensor & resize_(Tensor & self, IntArrayRef size) {`
	`164`	`+Tensor& resize_(`
	`165`	`+ Tensor& self,`
	`166`	`+ IntArrayRef size,`
	`167`	`+ c10::optional<MemoryFormat> optional_memory_format) {`
`165`	`168`	`auto& self_ = unpack(self, "self", 0);`
`166`	`169`	`if (as_variable_ref(self).requires_grad()) {`
`167`	`170`	`AT_ERROR("cannot resize variables that require grad");`
`@@ -173,7 +176,7 @@ Tensor & resize_(Tensor & self, IntArrayRef size) {`
`173`	`176`	`}`
`174`	`177`	`{`
`175`	`178`	`at::AutoNonVariableTypeMode non_var_type_mode(true);`
`176`		`- self_.resize_(size);`
	`179`	`+ self_.resize_(size, std::move(optional_memory_format));`
`177`	`180`	`}`
`178`	`181`	`return self;`
`179`	`182`	`}`