Avoid unnecessary copy in TensorSource

[lsy323]

Avoid at::Tensor copy in TensorSource if it's not necessary.

The copy operations are needed under 2 cases:

1. On XLA:GPU path, if tensor is on CUDA device, need to copy to CPU then pass to PJRT runtime to transfer to GPU. This is because @ysiraichi found passing CUDA tensor to PJRT runtime doesn't work, so doing the roundtrip as a workaround.
2. On XLA:TPU path, if the tensor is not contiguous, need to use the copy to make the memory contiguous. Because PJRT takes raw data ptr which expects data to be contiguous.

The copy operation needs to be blocking, since the transfer operation depends on the copied tensor.

[lsy323]

Hi @ysiraichi, just follow up on offline discussion on the copy operation. PTAL at the PR, thanks!

[ysiraichi]

As a side note, we can use the DLPack machinery for doing the CUDA to XLA:CUDA transfer (that wasn't implemented at the time I worked on this). I will open an issue for this.

[ysiraichi]

As far as I understand it, tensor.to(...) (without the copy argument) already checks whether it should actually copy or not. So, what do you think of reverting to the old tensor.to(...) usage, but removing the copy argument, instead?

[lsy323]

Hi @ysiraichi, I didn't find a tensor.to(...) without the copy arg in C++, is it only in python?

[ysiraichi]

You are right. But, I think we can just /* copy= */false.

[ysiraichi]

In the tensor.to(...) implementation, we can see that if copy=false, using the old implementation, that function will check whether we need to actually copy the tensor. No need to do it ourselves.

[pgmoka]

Left one small NIT. Otherwise, LGTM. If you could address the NIT before submission, that might be nice.

[pgmoka]

NIT: I personally prefer to have TODOs linked to issues, and then having those assigned to people. That way, things can be more easily followed-up if a contributor is no longer active

[lsy323]

@ysiraichi @tengyifei Actually took a 2nd thought on this. Skipping the copy seems to be unsafe if the underlying PJRT transfer is async.

def gen_rand_on_xla():
  t = torch.rand(100000)
  t_xla = t.to('xla') # Async transfer start
  return t_xla

my_rand_on_xla = gen_rand_on_xla()
# After the above line, async transfer may not be done, but CPU tensor is already out of scope and may be gc'ed.

[ysiraichi]

Hmm... I don't think this is the case. That's because PyTorch tensors are ref-counted in the C++ side (unless otherwise specified). So, if we hold a C++ Tensor type data, that won't be deleted.

[ysiraichi]

By the way, I still don't think you need the if-else there. I believe you can just leave the old tensor.to(...), while setting copy=false. PyTorch should verify whether a copy is needed for us.

[ysiraichi]

LGTM.

[lsy323]

By the way, I still don't think you need the if-else there. I believe you can just leave the old tensor.to(...), while setting copy=false. PyTorch should verify whether a copy is needed for us.

Thank you so much @ysiraichi for the suggestion! Updated accordingly.