2424#include " paddle/phi/api/include/api.h"
2525#include " paddle/phi/common/memory_utils.h"
2626#include " paddle/phi/core/ddim.h"
27+ #include " paddle/phi/core/memory/malloc.h"
2728
2829namespace at {
2930
@@ -49,11 +50,24 @@ inline int64_t ResizeCheckedNumel(at::IntArrayRef size) {
4950 return numel;
5051}
5152
53+ inline size_t ResizeCheckedStorageBytes (int64_t numel,
54+ size_t itemsize,
55+ size_t storage_offset_bytes) {
56+ const auto numel_size = static_cast <size_t >(numel);
57+ TORCH_CHECK (
58+ itemsize == 0 || numel_size <= (std::numeric_limits<size_t >::max () -
59+ storage_offset_bytes) /
60+ itemsize,
61+ " resize_ size is too large in bytes"
62+ return storage_offset_bytes + numel_size * itemsize;
63+ }
64+
5265} // namespace detail
5366
5467// resize_ - operate on the underlying DenseTensor directly so we preserve
55- // storage semantics across shrink/grow round-trips and only reallocate when
56- // the requested shape exceeds the current storage capacity.
68+ // storage semantics across shrink/grow round-trips. When growth exceeds the
69+ // current capacity, expand the shared storage itself so aliasing views keep
70+ // their storage offset and existing storage contents stay intact.
5771inline const at::Tensor& Tensor::resize_ (
5872 at::IntArrayRef size,
5973 ::std::optional<at::MemoryFormat> memory_format) const {
@@ -72,35 +86,33 @@ inline const at::Tensor& Tensor::resize_(
7286 " resize_ is not allowed on an undefined tensor"
7387
7488 const size_t itemsize = phi::SizeOf (dense_tensor->dtype ());
75- const size_t old_numel = static_cast <size_t >(tensor_.numel ());
76- const size_t new_numel_size = static_cast <size_t >(new_numel);
77- const size_t required_bytes = new_numel_size * itemsize;
78- const size_t available_bytes =
79- dense_tensor->Holder () == nullptr
80- ? 0
81- : dense_tensor->Holder ()->size () - dense_tensor->meta ().offset ;
89+ const size_t new_storage_bytes = detail::ResizeCheckedStorageBytes (
90+ new_numel, itemsize, dense_tensor->meta ().offset );
91+ const size_t current_storage_bytes =
92+ dense_tensor->Holder () == nullptr ? 0 : dense_tensor->Holder ()->size ();
8293
83- if (required_bytes <= available_bytes || new_numel == 0 ) {
94+ if (new_storage_bytes <= current_storage_bytes || new_numel == 0 ) {
8495 dense_tensor->Resize (dims);
8596 return *this ;
8697 }
8798
99+ // Sync through the compat Storage path first so the DenseTensor holder is a
100+ // live StorageHolderView backed by shared StorageImpl.
101+ auto storage = this ->storage ();
88102 const auto old_holder = dense_tensor->Holder ();
89103 TORCH_CHECK (old_holder != nullptr ,
90104 " resize_ cannot grow a tensor without allocated storage"
91- const size_t old_offset = dense_tensor->meta ().offset ;
92- const size_t copy_bytes = std::min (old_numel, new_numel_size) * itemsize;
93105 const phi::Place place = old_holder->place ();
94- const void * old_data =
95- old_holder == nullptr
96- ? nullptr
97- : reinterpret_cast <const uint8_t *>(old_holder->ptr ()) + old_offset;
98-
99- dense_tensor->ResizeAndAllocate (phi::make_ddim (dims));
100- void * new_data = dense_tensor->data ();
101- if (copy_bytes > 0 && old_data != nullptr && old_data != new_data) {
102- phi::memory_utils::Copy (place, new_data, place, old_data, copy_bytes);
106+ auto new_holder = paddle::memory::AllocShared (place, new_storage_bytes);
107+ TORCH_CHECK (new_holder != nullptr , " resize_ failed to allocate storage"
108+ const size_t copy_bytes = std::min (old_holder->size (), new_storage_bytes);
109+ if (copy_bytes > 0 && old_holder->ptr () != nullptr &&
110+ old_holder->ptr () != new_holder->ptr ()) {
111+ phi::memory_utils::Copy (
112+ place, new_holder->ptr (), place, old_holder->ptr (), copy_bytes);
103113 }
114+ storage.set_data_ptr_noswap (std::move (new_holder));
115+ dense_tensor->Resize (phi::make_ddim (dims));
104116 return *this ;
105117}
106118
0 commit comments