2022-12-12 nightly release (b95ea4f)

pytorch · ezyang · Sep 11, 2022 · Sep 12, 2022 · Sep 13, 2022 · Sep 14, 2022
commit 008d7ee6165e41e6ad4fa6b7b1cc687010ce725d
diff --git a/aten/src/ATen/EmptyTensor.cpp b/aten/src/ATen/EmptyTensor.cpp
@@ -106,13 +106,22 @@ size_t computeStorageNbytes(
 #endif
 }
 
+SymInt computeStorageNbytesContiguous(
+    SymIntArrayRef sizes,
+    const SymInt& itemsize_bytes,
+    const SymInt& storage_offset
+  ) {
+  const auto numel = c10::multiply_integers(sizes);
+  return itemsize_bytes * (storage_offset + numel);
+}
+
 // not including mobile-only macros in this function,
 // since mobile shouldn't be using symints.
 SymInt computeStorageNbytes(
     SymIntArrayRef sizes,
     SymIntArrayRef strides,
-    SymInt itemsize_bytes,
-    SymInt storage_offset
+    const SymInt& itemsize_bytes,
+    const SymInt& storage_offset
   ) {
   TORCH_CHECK(
     sizes.size() == strides.size(),
@@ -135,28 +144,28 @@ SymInt computeStorageNbytes(
   return itemsize_bytes * (storage_offset + size);
 }
 
-TensorBase empty_generic(
-    IntArrayRef size,
+template <typename T>
+TensorBase _empty_generic(
+    ArrayRef<T> size,
     c10::Allocator* allocator,
     c10::DispatchKeySet ks,
     ScalarType scalar_type,
     c10::optional<c10::MemoryFormat> memory_format_opt) {
   at::detail::check_size_nonnegative(size);
   at::detail::raise_warning_for_complex_half(scalar_type);
   caffe2::TypeMeta dtype = scalarTypeToTypeMeta(scalar_type);
-  size_t size_bytes = computeStorageNbytesContiguous(size, dtype.itemsize());
+  auto size_bytes = computeStorageNbytesContiguous(size, dtype.itemsize());
   auto storage_impl = c10::make_intrusive<StorageImpl>(
       c10::StorageImpl::use_byte_size_t(),
       size_bytes,
-      allocator->allocate(size_bytes),
       allocator,
       /*resizeable=*/true);
 
   auto tensor = detail::make_tensor_base<TensorImpl>(
       std::move(storage_impl), ks, dtype);
   // Default TensorImpl has size [0]
   if (size.size() != 1 || size[0] != 0) {
-    tensor.unsafeGetTensorImpl()->set_sizes_contiguous(size);
+    tensor.unsafeGetTensorImpl()->generic_set_sizes_contiguous(size);
   }
 
   if (memory_format_opt.has_value()) {
@@ -169,6 +178,15 @@ TensorBase empty_generic(
   return tensor;
 }
 
+TensorBase empty_generic(
+    IntArrayRef size,
+    c10::Allocator* allocator,
+    c10::DispatchKeySet ks,
+    ScalarType scalar_type,
+    c10::optional<c10::MemoryFormat> memory_format_opt) {
+  return _empty_generic(size, allocator, ks, scalar_type, memory_format_opt);
+}
+
 template <typename T>
 TensorBase _empty_strided_generic(
     T size,
@@ -338,59 +356,10 @@ TensorBase empty_symint_meta(
   c10::optional<bool> pin_memory_opt,
   c10::optional<c10::MemoryFormat> memory_format_opt
 ) {
-  auto device = device_or_default(device_opt);
-  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(device.type() == DeviceType::Meta);
-  // NB: because there is no SparseMeta (yet), non-strided layout is
-  // exerciseable
-  TORCH_CHECK_NOT_IMPLEMENTED(
-    layout_or_default(layout_opt) == Layout::Strided,
-    "non-strided meta tensors not supported yet"
-  );
-
-  auto scalar_type = dtype_or_default(dtype_opt);
   auto *allocator = GetAllocator(kMeta);
-  constexpr c10::DispatchKeySet meta_dks(c10::DispatchKey::Meta);
-  at::detail::check_size_nonnegative(size);
-  at::detail::raise_warning_for_complex_half(scalar_type);
-  caffe2::TypeMeta dtype = scalarTypeToTypeMeta(scalar_type);
-  SymInt size_bytes = dtype.itemsize();
-  for (auto s : size) {
-    size_bytes = size_bytes * s;
-  }
-  auto storage_impl = c10::make_intrusive<StorageImpl>(
-      c10::StorageImpl::use_byte_size_t(),
-      size_bytes,
-      allocator,
-      /*resizeable=*/true);
-
-  auto tensor = detail::make_tensor_base<TensorImpl>(
-      std::move(storage_impl), meta_dks, dtype);
-
-  int64_t dim = size.size();
-  std::vector<SymInt> strides;
-  strides.resize(dim);
-
-  // TODO: Move this into TensorImpl
-  auto memory_format = memory_format_opt.value_or(MemoryFormat::Contiguous);
-  switch (memory_format) {
-    case MemoryFormat::Contiguous: {
-      if (dim > 0) {
-        const auto last_idx = dim - 1;
-        strides.at(last_idx) = 1;
-        for (auto i = last_idx - 1; i >= 0; --i) {
-          // TODO: max with 1
-          strides.at(i) = strides.at(i+1) * size.at(i+1);
-        }
-      }
-      break;
-    }
-    default:
-      TORCH_CHECK(0, "other memory format not implemented yet");
-  }
-
-  tensor.unsafeGetTensorImpl()->set_sizes_and_strides(size, strides);
-
-  return tensor;
+  constexpr c10::DispatchKeySet ks(c10::DispatchKey::Meta);
+  auto scalar_type = dtype_or_default(dtype_opt);
+  return _empty_generic(size, allocator, ks, scalar_type, memory_format_opt);
 }
 
 TensorBase empty_meta(

diff --git a/aten/src/ATen/EmptyTensor.h b/aten/src/ATen/EmptyTensor.h
@@ -20,6 +20,10 @@ TORCH_API size_t computeStorageNbytesContiguous(
     IntArrayRef sizes,
     size_t itemsize,
     size_t storage_offset = 0);
+TORCH_API SymInt computeStorageNbytesContiguous(
+    SymIntArrayRef sizes,
+    const SymInt& itemsize,
+    const SymInt& storage_offset = 0);
 TORCH_API size_t computeStorageNbytes(
     IntArrayRef sizes,
     IntArrayRef strides,
@@ -28,8 +32,8 @@ TORCH_API size_t computeStorageNbytes(
 TORCH_API SymInt computeStorageNbytes(
     SymIntArrayRef sizes,
     SymIntArrayRef strides,
-    SymInt itemsize,
-    SymInt storage_offset = 0);
+    const SymInt& itemsize,
+    const SymInt& storage_offset = 0);
 
 TORCH_API TensorBase empty_generic(
     IntArrayRef size,

diff --git a/aten/src/ATen/native/Resize.cpp b/aten/src/ATen/native/Resize.cpp
@@ -10,20 +10,22 @@
 #else
 #include <ATen/ops/resize_as_native.h>
 #include <ATen/ops/resize_native.h>
+#include <ATen/ops/resize.h>
 #endif
 
 namespace at { namespace native {
 
 // Returns true if resize is necessary
-bool resize_output_check(const Tensor& output, IntArrayRef shape) {
+template <typename T>
+bool _resize_output_check(const Tensor& output, ArrayRef<T> shape) {
   // Tests for resizing of tensors with one or more elements
-  if (output.sizes().equals(shape)) {
+  if (at::symint::sizes<T>(output).equals(shape)) {
     return false;
   }
-  if (output.numel() != 0) {
+  if (at::symint::numel<T>(output) != 0) {
     TORCH_WARN(
       "An output with one or more elements was resized since it had ",
-      "shape ", output.sizes(), ", which does not match the required ",
+      "shape ", at::symint::sizes<T>(output), ", which does not match the required ",
       "output shape ", shape, ". ",
       "This behavior is deprecated, and in a future PyTorch release outputs ",
       "will not be resized unless they have zero elements. You can explicitly ",
@@ -33,25 +35,50 @@ bool resize_output_check(const Tensor& output, IntArrayRef shape) {
   return true;
 }
 
-bool resize_output(const Tensor& output, IntArrayRef shape) {
-  if (resize_output_check(output, shape)) {
+bool resize_output_check(const Tensor& output, IntArrayRef shape) {
+  return _resize_output_check(output, shape);
+}
+
+bool resize_output_check_symint(const Tensor& output, SymIntArrayRef shape) {
+  return _resize_output_check(output, shape);
+}
+
+void native_resize_(const Tensor& output, IntArrayRef shape) {
+  native::resize_(output, shape);
+}
+
+void native_resize_(const Tensor& output, SymIntArrayRef shape) {
+  native::resize__symint(output, shape);
+}
+
+template <typename T>
+bool _resize_output(const Tensor& output, ArrayRef<T> shape) {
+  if (_resize_output_check<T>(output, shape)) {
     // avoid a redispatch for cpu and cuda.
     // TODO: when resize_cuda_ is re-written to be unified with resize_,
     // we can provide the same benefit for cuda.
     //
     // TODO(#61485): functorch wrapped tensors should not go through the
     // fast path. This is a hack, longer term solutions are in the issue
     if (output.is_cpu() && !isTensorSubclassLike(output)) {
-      at::native::resize_(output, shape);
+      native_resize_(output, shape);
     } else {
-      output.resize_(shape);
+      at::symint::resize_<T>(output, shape);
     }
     return true;
   } else {
     return false;
   }
 }
 
+bool resize_output(const Tensor& output, IntArrayRef shape) {
+  return _resize_output(output, shape);
+}
+
+bool resize_output_symint(const Tensor& output, SymIntArrayRef shape) {
+  return _resize_output(output, shape);
+}
+
 const Tensor& _resize_output_(const Tensor& self, IntArrayRef shape, c10::Device device) {
   TORCH_CHECK(self.device() == device, "out Tensor doesn't have the correct device set");
   at::native::resize_output(self, shape);
@@ -126,16 +153,92 @@ const Tensor& resize_as_(
   return result;
 }
 
-const Tensor& resize_(
-    const Tensor& self,
+
+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));
+}
+
+static void maybe_resize_storage_meta(TensorImpl* self, c10::SymInt new_size_bytes) {
+  // It does not make sense to try to resize a storage
+  // to hold 0 elements, and this can break
+  // if storage_offset is positive but
+  // new_size is 0, so just bail in that case
+  // (same comment is in Resize.h)
+  if (self->sym_numel() == 0) {
+    return;
+  }
+
+  const Storage& storage = self->unsafe_storage();
+  if (!storage) {
+    TORCH_INTERNAL_ASSERT(0, "NYI, this should only be Caffe2");
+  } else if (new_size_bytes > storage.nbytes()) {
+    resize_bytes_meta(storage.unsafeGetStorageImpl(), std::move(new_size_bytes));
+  }
+}
+
+static void _maybe_resize_storage(TensorImpl* self, int64_t new_size_bytes) {
+  maybe_resize_storage_cpu(self, new_size_bytes);
+}
+
+static void _maybe_resize_storage(TensorImpl* self, c10::SymInt new_size_bytes) {
+  maybe_resize_storage_meta(self, std::move(new_size_bytes));
+}
+
+template <typename T>
+TensorImpl* _resize_impl_(
+    TensorImpl* self,
+    ArrayRef<T> size,
+    at::OptionalArrayRef<T> stride,
+    bool resize_storage) {
+  if (self->generic_sizes<T>() == size && (!stride || self->generic_strides<T>() == stride.value())) {
+    return self;
+  }
+
+  const auto itemsize = self->dtype().itemsize();
+  const auto storage_offset = self->generic_storage_offset<T>();
+  T storage_size = T(1);
+  if (stride) {
+    self->set_sizes_and_strides(size, *stride);
+    storage_size = at::detail::computeStorageNbytes(
+        size, *stride, itemsize, storage_offset);
+      } else {
+    self->generic_set_sizes_contiguous(size);
+    storage_size = at::detail::computeStorageNbytesContiguous(
+        size, itemsize, storage_offset);
+  }
+
+  if (resize_storage) {
+    _maybe_resize_storage(self, std::move(storage_size));
+  }
+
+  return self;
+}
+
+TensorImpl* resize_impl_cpu_(
+    TensorImpl* self,
     IntArrayRef size,
+    at::OptionalIntArrayRef stride,
+    bool resize_storage) {
+  return _resize_impl_(self, size, stride, resize_storage);
+}
+
+TensorImpl* resize_impl_meta_(
+    TensorImpl* self,
+    c10::SymIntArrayRef size,
+    at::OptionalSymIntArrayRef stride,
+    bool resize_storage = true) {
+  return _resize_impl_(self, size, stride, resize_storage);
+}
+
+template <typename T>
+const Tensor& _resize_(
+    const Tensor& self,
+    ArrayRef<T> size,
     c10::optional<MemoryFormat> optional_memory_format) {
-  if (self.has_names()) {
-    return resize_named_tensor_(self, size, optional_memory_format);
-  }
   auto* self_ = self.unsafeGetTensorImpl();
   // NOLINTNEXTLINE(bugprone-argument-comment)
-  resize_impl_cpu_(self_, size, /*strides=*/c10::nullopt);
+  _resize_impl_<T>(self_, size, /*strides=*/c10::nullopt, true);
   if (optional_memory_format.has_value()) {
     auto memory_format =
         optional_memory_format.value();
@@ -148,5 +251,23 @@ const Tensor& resize_(
   return self;
 }
 
+const Tensor& resize_(
+    const Tensor& self,
+    IntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
+  if (self.has_names()) {
+    return resize_named_tensor_(self, size, optional_memory_format);
+  }
+  return _resize_(self, size, optional_memory_format);
+}
+
+const Tensor& resize__symint(
+    const Tensor& self,
+    c10::SymIntArrayRef size,
+    c10::optional<MemoryFormat> optional_memory_format) {
+  TORCH_INTERNAL_ASSERT(!self.has_names())
+  return _resize_(self, size, optional_memory_format);
+}
+
 } // namespace native
 } // namespace at