pearu/pydlpack

how can i use this with my c extension?

twmht opened this issue · 3 comments

Hi,

I am writing a c extension using Python.h, and I have a pointer to gpu memory, I want to return a "Tensor" representing the data to python.

What is the recommanation way to do that?

The answer to your questions depends on what array library object (say, pytorch, or cupy, or jax, etc) you wish to return from your extension module.

pydlpack is for exchanging device data between array libraries from Python. So, the raised issue is OT to this project.

@pearu

My question is, if I can already return CuPy, which inherently implements DLpack and can exchange data with different frameworks, in what scenarios would I still need to use PyDLPack?

My question is, if I can already return CuPy, which inherently implements DLpack and can exchange data with different frameworks, in what scenarios would I still need to use PyDLPack?

This is an excellent question, thank you!

For DLPack protocol to be effective, both provider and consumer libraries need to implement the DLPack protocol: the provider library must implement __dlpack__ method on its array object and the consumer library must implement from_dlpack function.

PyDLPack comes into play when a provider library does not implement DLPack protocol but it implements some other data exchange protocol such as buffer protocol or array interface protocol or CUDA array interface protocol. PyDLPack will then construct the DLPack interface in runtime for provider libraries that otherwise lack the export DLPack support.

Here's a (incomplete) list of objects that PyDLPack is able to make DLPack-aware that otherwise are not: Numba DeviceNDArray, Python bytes, Python bytearray, Python array.

Another use case where PyDLPack can be handy is to exchange data between immutable and mutable array objects. For instance, JAX exports its arrays with readonly flag set to True that PyTorch is not able to import because it does not support readonly tensors. Using PyDLPack one can overwrite the readonly flag and thus enabling importing JAX arrays to PyTorch (this assumes that the corresponding PyTorch tensors are treated as readonly, of course).