def _test_wait_stream(self, source, target, cuda_sleep=None): with use_stream(target): if is_cuda(target): cuda_sleep(0.5) x = torch.ones(100, 100, device=get_device(target)) wait_stream(source, target) with use_stream(source): assert x.sum().item() == 10000
def backward(ctx: Context, *grad_output: Tensor, ) -> Tuple[Optional[Tensor], ...]: prev_stream = ctx.prev_stream next_stream = ctx.next_stream grad_input: Deque[Tensor] = deque(maxlen=len(grad_output)) input_stream = current_stream(get_device(prev_stream)) with use_stream(prev_stream), use_stream(next_stream): for x in reversed(grad_output): y = x.to(get_device(prev_stream)) grad_input.appendleft(y) # 'next_stream' is not where 'x' has been allocated. record_stream(x, next_stream) # 'y' has been allocated on 'prev_stream'. # It might be used on the current stream captured as 'input_stream'. record_stream(y, input_stream) grad_streams: Tuple[Optional[Tensor], ...] = (None, None) return grad_streams + tuple(grad_input)
def forward(ctx: Context, # type: ignore prev_stream: AbstractStream, next_stream: AbstractStream, *input: Tensor, ) -> Tensors: ctx.prev_stream = prev_stream ctx.next_stream = next_stream output = [] output_stream = current_stream(get_device(next_stream)) with use_stream(prev_stream), use_stream(next_stream): for x in input: y = x.to(get_device(next_stream)) output.append(y) # 'prev_stream' is not where 'x' has been allocated. record_stream(x, prev_stream) # 'y' has been allocated on 'next_stream'. # It might be used on the current stream captured as 'output_stream'. record_stream(y, output_stream) return tuple(output)
def forward( ctx: Context, # type: ignore portal: Portal, prev_stream: AbstractStream, next_stream: AbstractStream, phony: Tensor, ) -> Tensor: ctx.portal = portal assert portal.tensor is not None portal.tensor, = Copy.forward(ctx, prev_stream, next_stream, portal.tensor) phony = get_phony(get_device(next_stream), requires_grad=False) return phony.detach()
def _test_copy_wait(prev_stream, next_stream, cuda_sleep=None): device = get_device(prev_stream) with use_stream(prev_stream): if is_cuda(prev_stream): cuda_sleep(0.5) x = torch.ones(100, device=device, requires_grad=True) y, = Copy.apply(prev_stream, next_stream, x) y, = Wait.apply(prev_stream, next_stream, x) with use_stream(next_stream): assert torch.allclose(y.sum(), torch.tensor(100.0, device=device)) y.norm().backward() with use_stream(prev_stream): assert torch.allclose(x.grad.sum(), torch.tensor(10.0, device=device))
def test_get_device_cuda(self): stream = current_stream(torch.device('cuda')) assert get_device(stream).type == 'cuda'
def test_get_device_cpu(self): assert get_device(CPUStream).type == 'cpu'