wjakob / nanobind

nanobind: tiny and efficient C++/Python bindings
BSD 3-Clause "New" or "Revised" License
2.29k stars 191 forks source link

Add cupy support #594

Closed JoOkuma closed 4 months ago

JoOkuma commented 4 months ago

Hi, this PR adds cupy support.

What additional tests do you recommend adding?

Do you think it would make sense for np.ndarray with nb::device::cuda be considered cp.ndarray automatically?

Reference: #510

wjakob commented 4 months ago

This looks great. It's tricky to test these CUDA frameworks on GitHub actions. If it works for you, it is good enough for me. Please list the CuPy enumeration entry in the docs (both the ndarray overview file and the api_extra.rst file), then it's good to go.

JoOkuma commented 4 months ago

Thanks, @wjakob; I updated the docs.

I'm getting an error on the cupy side in some environments, lets wait for me to understand this before merging this PR.

E           RuntimeError: nanobind::detail::ndarray_wrap(): could not import ndarray: Traceback (most recent call last):
E             File "cupy/_core/dlpack.pyx", line 186, in cupy._core.dlpack.DLPackMemory.__init__
E           RuntimeError: CuPy is built against CUDA, different from the backend that backs the incoming DLPack tensor
wjakob commented 4 months ago

Probably an attempt to convert a CPU dlpack tensor into a CuPy tensor. Perhaps they don't implement the copy back to the GPU.

JoOkuma commented 4 months ago

You're right. It doesn't support a CPU dlpack.

And the error was on my side. I did not set the nb::device::cuda::value to the returning cuda array. I think it's ready to merge.

Thanks!

wjakob commented 4 months ago

And the error was on my side. I did not set the nb::device::cuda::value to the returning cuda array. I think it's ready to merge.

Does this have any implications for the nanobind PR? I did not see further commits/amends since you raised this issue.

JoOkuma commented 4 months ago

No, @wjakob. It worked once I fixed how I was using nanobind in my other codebase.

wjakob commented 4 months ago

Got it. Thanks for the PR!