Support resize on meta storage (#101988)

Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #101988 Approved by: https://github.com/albanD, https://github.com/bdhirsh
pytorch · May 23, 2023 · 7d1ba0a · 7d1ba0a · pytorchmergebot · May 24, 2023
1 parent 51ff408
commit 7d1ba0a
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 4 deletions.
diff --git a/aten/src/ATen/native/Resize.cpp b/aten/src/ATen/native/Resize.cpp
@@ -156,7 +156,7 @@ const Tensor& resize_as_(
 }
 
 
-static void resize_bytes_meta(StorageImpl* storage, c10::SymInt size_bytes) {
+void resize_bytes_meta(StorageImpl* storage, c10::SymInt size_bytes) {
   TORCH_CHECK(storage->resizable(), "Trying to resize storage that is not resizable");
   storage->set_nbytes(std::move(size_bytes));
 }

diff --git a/aten/src/ATen/native/Resize.h b/aten/src/ATen/native/Resize.h
@@ -34,6 +34,7 @@ TORCH_API bool resize_output_check(const Tensor& output, IntArrayRef shape);
 TORCH_API bool resize_output_check_symint(const Tensor& output, SymIntArrayRef shape);
 
 TORCH_API void resize_bytes_cpu(StorageImpl* storage, size_t size_bytes);
+TORCH_API void resize_bytes_meta(StorageImpl* storage, c10::SymInt size_bytes);
 
 static inline void maybe_resize_storage_cpu(TensorImpl* self, size_t new_size_bytes) {
   // It does not make sense to try to resize a storage

diff --git a/test/test_torch.py b/test/test_torch.py
@@ -350,9 +350,6 @@ def test_storage_meta_errors(self, device, dtype):
             with self.assertRaisesRegex(NotImplementedError, r'Cannot copy out'):
                 s0.pin_memory()
 
-        with self.assertRaisesRegex(RuntimeError, r'got unexpected device type'):
-            s0.resize_(10)
-
         with self.assertRaisesRegex(RuntimeError, r'only available on CPU'):
             s0.share_memory_()
 
@@ -369,6 +366,14 @@ def test_storage_meta_errors(self, device, dtype):
             with self.assertRaisesRegex(NotImplementedError, r'Cannot copy out'):
                 s1.copy_(s0)
 
+    @onlyCPU
+    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
+    def test_storage_meta_ok(self, device, dtype):
+        s0 = torch.TypedStorage([1, 2, 3, 4], device='meta', dtype=dtype)
+
+        # This is OK, it changes the meta storage size without allocating
+        s0.resize_(10)
+
     @onlyCUDA
     def test_module_share_memory(self):
         # Test fix for issue #80733

diff --git a/torch/csrc/StorageMethods.cpp b/torch/csrc/StorageMethods.cpp
@@ -128,6 +128,8 @@ static PyObject* THPStorage_resize_(PyObject* self, PyObject* number_arg) {
     const auto size_bytes = static_cast<size_t>(size_bytes_i);
     at::native::resize_bytes_cuda(storage.unsafeGetStorageImpl(), size_bytes);
 #endif
+  } else if (device_type == at::kMeta) {
+    at::native::resize_bytes_meta(storage.unsafeGetStorageImpl(), newsize);
   } else if (device_type == at::kPrivateUse1) {
     ptrdiff_t size_bytes_i = newsize;
     TORCH_CHECK(