frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix storage use_count (#157694)

Browse Source 
# Motivation
https://github.com/pytorch/pytorch/pull/155451 decoupled `torch._C._storage_Use_Count` from CUDA and introduced a corresponding unit test:
815545f2dd/test/test_torch.py (L257-L262)
However, this test fails when PyTorch is built with debug assertions enabled. @clee2000 disabled this UT in https://github.com/pytorch/pytorch/pull/156731. The root cause is that `_cdata` is obtained from an `intrusive_ptr`, not a `weak_intrusive_ptr`. As a result, calling `c10::weak_intrusive_ptr::use_count` on it triggers the internal assertion:
815545f2dd/c10/util/intrusive_ptr.h (L912-L917)
For example:
```python
a = torch.randn(10, device=device) # refcount=1, weakcount=1
prev_cf = torch._C._storage_Use_Count(a.untyped_storage()._cdata) # violate the assertation
```
This violates the expected invariant inside `weak_intrusive_ptr::use_count`, which assumes the pointer was originally constructed from a valid `weak_intrusive_ptr`. Actually, `storage_impl` is obtained from an `intrusive_ptr`.
815545f2dd/torch/csrc/Module.cpp (L2105-L2109)

# Solution
Use `c10::intrusive_ptr::use_count` instead.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/157694
Approved by: https://github.com/albanD
This commit is contained in:
Yu, Guangye
2025-07-07 13:37:04 +00:00
committed by PyTorch MergeBot
parent 8186af5a26
commit 1b58e7adab
2 changed files with 1 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/test_torch.py
View File

@ -250,10 +250,6 @@ class TestTorchDeviceType(TestCase):
@skipIfTorchDynamo("Not a suitable test for TorchDynamo")
@onlyNativeDeviceTypes
@unittest.skipIf(
"RelWithAssert" in torch.__config__.show(),
"failing in debug build, see https://github.com/pytorch/pytorch/pull/156731 for example",
)
def test_storage_use_count(self, device):
a = torch.randn(10, device=device)
prev_cf = torch._C._storage_Use_Count(a.untyped_storage()._cdata)

2
torch/csrc/Module.cpp
View File

@ -2105,7 +2105,7 @@ Call this whenever a new thread is created in order to propagate values from
py_module.def("_storage_Use_Count", [](size_t storage_impl_ptr) {
// NOLINTNEXTLINE(performance-no-int-to-ptr)
c10::StorageImpl* storage_impl = (c10::StorageImpl*)storage_impl_ptr;
return c10::raw::weak_intrusive_ptr::use_count(storage_impl);
return c10::raw::intrusive_ptr::use_count(storage_impl);
});
ASSERT_TRUE(