Example #1
0
    def test_record_stream_cuda(self, cuda_sleep):
        # This test detects unexpected block reallocation. For reliable test,
        # the stream to allocate tensors is isolated. The allocator will not
        # reuse free blocks which were allocated from another stream.
        stream_alloc = new_stream(torch.device("cuda"))
        with torch.cuda.stream(stream_alloc):
            x = torch.rand(1, device=torch.device("cuda"))

        stream = new_stream(torch.device("cuda"))
        record_stream(x, stream)
        with use_stream(stream):
            cuda_sleep(0.5)

        # 'x' is deleted at Python's perspective. But the block of 'x' is still
        # required for 'stream'. 'y' shouldn't be allocated to the block.
        data_ptr = x.data_ptr()
        del x
        stream_alloc.synchronize()
        with torch.cuda.stream(stream_alloc):
            y = torch.rand(1, device=torch.device("cuda"))
        assert y.data_ptr() != data_ptr

        # Pause Python until 'stream' finishes tasks queued. Now the block of
        # 'x' is free to be reallocated.
        wait_stream(CPUStream, stream)
        with torch.cuda.stream(stream_alloc):
            z = torch.rand(1, device=torch.device("cuda"))
        assert z.data_ptr() == data_ptr
Example #2
0
    def test_record_stream_shifted_view(self, cuda_sleep):
        # Issue: https://github.com/pytorch/pytorch/issues/27366
        stream_alloc = new_stream(torch.device("cuda"))
        with torch.cuda.stream(stream_alloc):
            x = torch.rand(2, device=torch.device("cuda"))

        y = x[1:]
        assert y.data_ptr() > x.data_ptr()

        stream = new_stream(torch.device("cuda"))
        with use_stream(stream):
            cuda_sleep(0.5)
        record_stream(y, stream)

        data_ptr = x.data_ptr()
        del x, y

        stream_alloc.synchronize()
        with torch.cuda.stream(stream_alloc):
            z = torch.rand(2, device=torch.device("cuda"))
        assert z.data_ptr() != data_ptr
Example #3
0
def test_copy_wait_cuda_cuda(cuda_sleep):
    prev_stream = current_stream(torch.device("cuda"))
    next_stream = new_stream(torch.device("cuda"))
    _test_copy_wait(prev_stream, next_stream, cuda_sleep)
Example #4
0
 def test_use_stream_cuda(self):
     stream = new_stream(torch.device("cuda"))
     with use_stream(stream):
         assert current_stream(torch.device("cuda")) == stream
Example #5
0
 def test_new_stream_cuda(self):
     stream = new_stream(torch.device("cuda"))
     assert isinstance(stream, torch.cuda.Stream)
     assert stream != torch.cuda.default_stream()
Example #6
0
 def test_new_stream_cpu(self):
     stream = new_stream(torch.device("cpu"))
     assert stream is CPUStream
Example #7
0
 def test_wait_stream_cuda_cuda(self, cuda_sleep):
     source = current_stream(torch.device("cuda"))
     target = new_stream(torch.device("cuda"))
     self._test_wait_stream(source, target, cuda_sleep)
Example #8
0
 def test_wait_stream_cuda_cpu(self, cuda_sleep):
     source = new_stream(torch.device("cuda"))
     target = CPUStream
     self._test_wait_stream(source, target, cuda_sleep)