Always synchronize on src and dst current streams when copying tensors #16966
Conversation
test/test_cuda.py
Outdated
| _test_copy(self, x0, x1, torch.zeros(5, 5, device=d1)) | ||
|
|
||
| x2 = torch.zeros(5, 5, device=d0) | ||
| _test_copy(self, x0, x2, torch.ones(5, 5, device=d1)) |
There was a problem hiding this comment.
This behavior is expected (as we do not sync if src and dst are the same), but looks weird to me. Should we do the sync using dst default stream regardless of whether src and dst are different?
There was a problem hiding this comment.
I don't understand. Which part is weird or unexpected here?
There was a problem hiding this comment.
These four lines, both are copying a zero tensor to x0, but the result differs, depending on whether the src and dst are on the same device.
x1 = torch.zeros(5, 5, device=d1)
_test_copy(self, x0, x1, torch.zeros(5, 5, device=d1))
x2 = torch.zeros(5, 5, device=d0)
_test_copy(self, x0, x2, torch.ones(5, 5, device=d1))There was a problem hiding this comment.
To make the multi-device case like the single-device case you would have to set the streams for both d0 and d1 (and synchronize with the current stream instead of default stream). Something like:
def _test_copy(self, x, y, output):
x_plus_one = x + 1
s0 = torch.cuda.Stream()
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream(device=y.device)
s3 = torch.cuda.Stream(device=y.device)
with torch.cuda.stream(s2), torch.cuda.stream(s0):
torch.cuda._sleep(50000000)
y.copy_(x_plus_one)
with torch.cuda.stream(s3), torch.cuda.stream(s1):
y.copy_(x)
s0.synchronize()
s1.synchronize()
|
The copy is still synchronizing on the default stream on the destination. It should be synchronizing on the current stream on the destination. That way it appears as-if the copy takes place in both the streams of the source and destination even though the kernel only runs on the current stream for the source's device. pytorch/aten/src/ATen/native/cuda/Copy.cu Line 78 in c5be4c5 and pytorch/aten/src/ATen/native/cuda/Copy.cu Line 153 in c5be4c5 |
|
rerun tests after merging #17439 |
mrshenli
changed the title
There was a problem hiding this comment.
@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
mrshenli
changed the title
Summary: fixes #15568 Pull Request resolved: pytorch/pytorch#16966 Differential Revision: D14213144 Pulled By: mrshenli fbshipit-source-id: 2fcf5e07895fde80b4aee72e2736b0def876d21f
4 participants
fixes #15568