[CUDA] Remove footgun related to non-blocking copies

[embg]

Edit: I filed this issue because I believed the following code has a use-after-free.

def foo():
    A = torch.rand(SIZE, device="cpu", pin_memory=True)
    B = A.cuda(non_blocking=True)
    return B

C = foo()
# do stuff with C

It turns out that PyTorch actually keeps the pinned memory alive until the memcpy event is complete. There is no bug. Pretty cool!

[colesbury]

I don't think this should result in a use after free. Are you sure that's the case?

pytorch/aten/src/ATen/native/cuda/Copy.cu

Line 381 in 54a932b

CachingHostAllocator_recordEvent(ptr, ctx, stream);

The above call keeps the pinned memory alive until the copy finishes.

You don't want to use cuStreamAddCallback() -- it does some nasty synchronization if I remember correctly.

[colesbury]

Here's a little example program that demonstrates that the pinned memory tensor does not get reused until the copy completes. You can add a torch.cuda._sleep() to delay the non-blocking copy kernel.

import torch
import time


def get_cycles_per_ns():
    cycles = 100000000
    torch.cuda.synchronize()
    start = time.perf_counter()
    torch.cuda._sleep(cycles)
    torch.cuda.synchronize()
    end = time.perf_counter()
    delta = (end - start) * 1e9
    return cycles / delta


CYCLES_PER_NS = get_cycles_per_ns()
SIZE = 1000

def cuda_sleep(ns):
    torch.cuda._sleep(int(ns * CYCLES_PER_NS))


def main():
    x = torch.randn(10, device='cuda')

    print("Ready")
    A = torch.ones((SIZE,), pin_memory=True)
    ptr1 = A.data_ptr()

    # Delay the non-blocking copy behind a kernel that takes ~5 seconds
    cuda_sleep(5 * 1000 * 1000 * 1000)

    B = A.cuda(non_blocking=True)

    del A
    C = torch.empty(SIZE, pin_memory=True)
    C.fill_(2)
    ptr2 = C.data_ptr()
    print("filled")

    total = B.sum().item()
    print(f"{total=}")
    print(f"{ptr1=} {ptr2=}")

    assert total == SIZE


if __name__ == '__main__':
    main()

[embg]

@colesbury You're right, thanks for the code pointer! (My original issue was not based on a real-life example, I just didn't read the code properly).

[embg]

@colesbury I have one more question for you. If the lifetimes are all managed by PyTorch, why is non_blocking=True not the default behavior whenever a CPU tensor uses pinned memory?

[colesbury]

Probably a mix of historical reasons, backwards compatibility, and concern about surprising behaviors. For example, there's still the hazard of:

def foo():
    A = torch.ones(SIZE, device="cpu", pin_memory=True)
    B = A.cuda(non_blocking=True)
    A.fill_(2)  # BUG! may overlap with non-blocking copy
    return B

assert foo().sum() == SIZE

Although deallocation of A should be safe, modifying A in-place can mess up the copy. This is different from dealing with CUDA tensors, where everything is happening in stream-order. It also different from default (blocking) CPU-GPU copies, where the synchronization prevents this hazard.

I think we could handle this hazard in PyTorch as well, but it would be a more invasive change.