Skip to content

Commit 9069aec

Browse files
youge325Copilot
andauthored
[Cpp API Compatibility] Fix resize_ with storage alignment (PaddlePaddle#78609)
--------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent d7543b4 commit 9069aec

3 files changed

Lines changed: 66 additions & 21 deletions

File tree

‎paddle/phi/api/include/compat/ATen/ops/resize.h‎

Lines changed: 33 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
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

2829
namespace 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.
5771
inline 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

‎paddle/phi/api/include/compat/ATen/ops/slice.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ inline at::Tensor slice(const at::Tensor& self,
2323
::std::optional<int64_t> start = ::std::nullopt,
2424
::std::optional<int64_t> end = ::std::nullopt,
2525
int64_t step = 1) {
26+
// Materialize the compat StorageHolderView before creating the slice so the
27+
// base tensor and its views observe the same shared storage during resize_.
28+
(void)self.storage();
2629
return paddle::experimental::slice(
2730
self._PD_GetInner(),
2831
{dim},

‎test/cpp/compat/ATen_resize_test.cc‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
#include <ATen/Functions.h>
1616
#include <ATen/core/TensorBody.h>
17+
#include <ATen/ops/as_strided.h>
1718
#include <ATen/ops/resize.h>
1819
#include <ATen/ops/tensor.h>
1920
#include <c10/core/ScalarType.h>
@@ -256,3 +257,32 @@ TEST(TensorResizeTest, ResizeChain) {
256257
ASSERT_EQ(t.sizes()[0], 3);
257258
ASSERT_EQ(t.sizes()[1], 4);
258259
}
260+
261+
// Test that resizing a view with shared storage copies data from the start of
262+
// the storage, not from the view's offset, to preserve the original data and
263+
// storage semantics.
264+
TEST(TensorResizeTest, ResizeSliceSharedStorageCopiesFromStorageStart) {
265+
// ta = [1, 2, 3, 4], tb = [2, 3, 4]
266+
// Build tb through as_strided so it is a view with a non-zero storage
267+
// offset even when backend slice kernels materialize copies.
268+
at::Tensor ta = at::tensor({1, 2, 3, 4}, at::kInt);
269+
at::Tensor tb = ta.as_strided({3}, {1}, 1);
270+
271+
tb.resize_(4);
272+
273+
// After resize, tb[0] and ta[1] must point to the exact same address.
274+
ASSERT_EQ(tb.data_ptr<int>(), ta.data_ptr<int>() + 1);
275+
276+
// The original storage contents should remain unchanged.
277+
ASSERT_EQ(ta[0].item<int>(), 1);
278+
ASSERT_EQ(ta[1].item<int>(), 2);
279+
ASSERT_EQ(ta[2].item<int>(), 3);
280+
ASSERT_EQ(ta[3].item<int>(), 4);
281+
282+
// PyTorch only preserves the pre-existing prefix here. The newly exposed
283+
// tail element after resize_ is uninitialized and should not be asserted.
284+
ASSERT_EQ(tb[0].item<int>(), 2);
285+
ASSERT_EQ(tb[1].item<int>(), 3);
286+
ASSERT_EQ(tb[2].item<int>(), 4);
287+
ASSERT_EQ(tb.numel(), 4);
288+
}

0 commit comments

Comments
 (0)