Auto-convert GPU arrays that support the __cuda_array_interface__ protocol #20584
Conversation
pytorchbot
added
module: internals
li-roy
added
the
triaged
ezyang
left a comment
ezyang
left a comment
There was a problem hiding this comment.
Logically, the code looks great! However, I would like the memory leaks to be fixed before merge. If you insist, just manually inserting the necessary decrefs is acceptable, however, I think using an RAII class from pybind11 will be much safer.
|
Thanks, I will use pybind11 to fix the memory leaks and also address the other issues. |
Now we only use `PyArray_DescrConverter()`
torch/csrc/utils/tensor_numpy.cpp
Outdated
| if(!PyTuple_Check(py_data) || PyTuple_GET_SIZE(py_data) != 2) { | ||
| throw TypeError("`data` must be a 2-tuple of (int, bool)"); | ||
| } | ||
| PyTuple_GET_ITEM(py_data, 0); |
There was a problem hiding this comment.
I don't think you intended to have this line?
ezyang
left a comment
ezyang
left a comment
There was a problem hiding this comment.
Will merge when tests pass.
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@ezyang is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
FYI @seibert, thought you'd like to know that PyTorch now supports |
7 participants
This PR implements auto-conversion of GPU arrays that support the
__cuda_array_interface__protocol (fixes #15601).If an object exposes the
__cuda_array_interface__attribute,touch.as_tensor()andtouch.tensor()will use the exposed device memory.Zero-copy
When using
touch.as_tensor(...,device=D)whereDis the same device as the one used in__cuda_array_interface__.Implicit copy
When using
touch.as_tensor(...,device=D)whereDis the CPU or another non-CUDA device.Explicit copy
When using
torch.tensor().Exception
When using
touch.as_tensor(...,device=D)whereDis a CUDA device not used in__cuda_array_interface__.Lifetime
torch.as_tensor(obj)tensor grabs a reference toobjso that the lifetime ofobjexceeds the tensor