record_stream() for shifted view tensors (#27371)
Summary:
Issue: https://github.com/pytorch/pytorch/issues/27366
The address of a view tensor might be shifted from the head of the storage.
```python
>>> x = torch.rand(10, 10, device=0, requires_grad=True)
>>> y = x[2:]
>>> hex(x.data_ptr())
'0x7f1b15c00000'
>>> hex(y.data_ptr())
'0x7f1b15c00050'
```
Currently, `Tensor.record_stream()` silently ignores shifted view tensors, because `CUDACachingAllocator` cannot find the block from the shifted address.
```c++
void recordStream(void* ptr, cuda::CUDAStream stream)
{
if (ptr) {
std::lock_guard<std::recursive_mutex> lock(mutex);
Block* block = find_allocated_block(ptr);
if (block) {
...
}
// 'block' is nullptr if 'ptr' is shifted.
}
}
```
So we cannot protect shifted view tensor which is used to compute or copy in an arbitrary stream against unexpected reallocation. Once we call `record_stream()` on a tensor, our intention is to protect the storage behind the tensor against reallocation until all works in the stream finish. This rule should be consistent regardless of the type of tensors including the view.
We can retrieve the head of the address from any types of tensors by `tensor.storage().data_ptr()`. Hence, I've thought it's better to pass to `recordStream()` rather than `tensor.data_ptr()` for consistent behavior.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/27371
Reviewed By: ezyang
Differential Revision: D17768558
Pulled By: albanD
fbshipit-source-id: 7705f52b0177625168edb6f71c07a029df471bc5
