[ENH]: Add DLPack casters to pybind11 #3858

Skylion007 · 2022-04-11T16:26:12Z

Required prerequisites

Make sure you've read the documentation. Your issue may be addressed there.
Search the issue tracker and Discussions to verify that this hasn't already been reported. +1 or comment there if it has.
Consider asking first in the Gitter chat room or in a Discussion.

Problem description

Recently, nanobind added a new type of array caster that uses DLPack to convert to and from array. We would like to backport this caster type to pybind. DLPack is a protocol for quickly and efficiently sharing arrays between programs that are backed by either CPU or GPU memory. For example, with a properly written caster one could transfer between a CUDA GPU Array to a PyTorch Tensor. While this is currently possible with custom caster in PyBind11, it would be nice to have official caster support to make this type of generic array sharing easier. Recent version of Numpy also have native APIs to create Numpy arrays from DLPack.

This is not a replacement for the Eigen / Numpy bindings, but a supplement to them. Our current bindings support versions of Numpy that are fairly old. Additionally, we have a lot of convenience functions for dealing with numpy arrays that DLPack likely will not support (reshaping, resizing etc.
This would be implemented as custom type casters. To use them, the user would need to specify an optional include (similar to how they already need to for Eigen/Numpy Bindings.
The testing suite and much of the code can be ripped from nanobind. Note that there are signficant differences between the type casters though: https://github.com/wjakob/nanobind#api-differences . We should ensure that the ported casters handle errors the pybind11 way.

Skylion007 · 2022-04-14T17:27:25Z

@oleksandr-pavlyk Are you interested in pursuing this issue? Given the text in PR #3866 seems like you already are working on an implementation, it would be nice to have official DLPack casters supported in pybind11.

oleksandr-pavlyk · 2022-04-14T17:55:24Z

I would love to.

PeterDykas · 2022-05-25T16:13:25Z

Sorry if I am reading this wrong but it seems like there may have been support added for DLPack but I was struggling to find information on it in the docs. Was the support added for DLPack tensors and is there anywhere you could point me to for finding out more?

oleksandr-pavlyk · 2022-05-25T16:28:45Z

@PeterDykas The support has been added in nanobind. This issue is to add support for DLPack casters for pybind11 following suite.

PeterDykas · 2022-05-25T16:56:38Z

Awesome thanks @oleksandr-pavlyk. When this is added it will be very useful! I would love to keep using pybind, is there a workaround that you know of to return DLPack tensors from pybind?

Tabrizian · 2022-07-06T03:27:02Z

@oleksandr-pavlyk @Skylion007 I would also be interested in working on this if it sounds good to you.

galv · 2022-07-11T22:03:06Z

FWIW I did create a custom type caster here: https://github.com/nvidia-riva/riva-asrlib-decoder/blob/main/include/riva/asrlib/decoder/pybind11_dlpack_caster.h I am going through legal approval to license it uner pybind11's BSD 3-clause license to I can open a PR.

One immediate observation is that it's quite hard to use DLManagedTensor directly, which is why you would presumably want the nb::tensor abstraction in pybind11.

(Note that I am a colleague of @Tabrizian)

Previously, this code would error out if the destructor happened to be a nullptr. This is incorrect. nullptrs are allowed for capsule destructors. "It is legal for a capsule to have a NULL destructor. This makes a NULL return code somewhat ambiguous; use PyCapsule_IsValid() or PyErr_Occurred() to disambiguate." See: https://docs.python.org/3/c-api/capsule.html#c.PyCapsule_GetDestructor I noticed this while working on a type caster related to pybind#3858 DLPack happens to allow the destructor not to be defined on a capsule, and I encountered such a case. See: https://github.com/dmlc/dlpack/blob/e2bdd3bee8cb6501558042633fa59144cc8b7f5f/include/dlpack/dlpack.h#L219

steven-johnson · 2022-10-07T17:32:04Z

Where does this task stand? Halide would very much like to support DLPack in our Python bindings, but nanobind isn't (yet) a good option for us.

#4221) * fix: PyCapsule_GetDestructor is allowed to return a nullptr destructor Previously, this code would error out if the destructor happened to be a nullptr. This is incorrect. nullptrs are allowed for capsule destructors. "It is legal for a capsule to have a NULL destructor. This makes a NULL return code somewhat ambiguous; use PyCapsule_IsValid() or PyErr_Occurred() to disambiguate." See: https://docs.python.org/3/c-api/capsule.html#c.PyCapsule_GetDestructor I noticed this while working on a type caster related to #3858 DLPack happens to allow the destructor not to be defined on a capsule, and I encountered such a case. See: https://github.com/dmlc/dlpack/blob/e2bdd3bee8cb6501558042633fa59144cc8b7f5f/include/dlpack/dlpack.h#L219 * Add test for the fix. * Update tests/test_pytypes.cpp I tried this locally and it works! I never knew that there are cases where `reinterpret_cast` does not work but `static_cast` does. Let's see if all compilers are happy with this. Co-authored-by: Aaron Gokaslan <[email protected]> * style: pre-commit fixes Co-authored-by: Ralf W. Grosse-Kunstleve <[email protected]> Co-authored-by: Ralf W. Grosse-Kunstleve <[email protected]> Co-authored-by: Aaron Gokaslan <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

galv · 2022-10-07T22:46:29Z

Hi @steven-johnson , I am using dlpack for a few things. The first thing was bindings for a CUDA-based beam search decoder for non-autoregressive CTC speech recognition models, for which I wrote a simple type caster here: https://github.com/nvidia-riva/riva-asrlib-decoder/blob/main/include/riva/asrlib/decoder/pybind11_dlpack_caster.h

I noticed a small error in it yesterday while working on a (not-yet-publicly-available) beam search decoder for auto-regressive speech recognition models, like RNN-T and attention-based models. These are more complicated because the decoder must call the neural network model multiple times during decoding, as opposed to just accepting a single output from a neural network model at the beginning. I noticed some errors in my original caster, as well as in pybind11, which prompted the PR here: 7c6f2f8

If you are okay using DLManagedTensor directly, instead of nanobind's batteries-included nb::tensor, take a look at this: https://gist.github.com/galv/4957a972587fbef28e64aeb6b03579ca It is my latest version of this.

The main problem with getting dlpack support in pybind11 is that, if you wanted something like nb::tensor, that is way, way, way more work than what I did. (Trust me, I looked into it by trying to implement it. It would end up becoming a very large undertaking if the powers-that-be wanted something like nb::tensor to consider this issue complete). If you want your C++ bindings to accept buffers from anything implementing the dlpack protocol, this type caster should work, modulo undiscovered bugs.

steven-johnson · 2022-10-07T22:52:25Z

if you wanted something like nb::tensor

I have only taken a quick skim into nanobind, so it's not clear to me what all is in nb::tensor at this point. I kinda assumed it is something like a numpy ndarray except also with support for the dlpack protocol (in addition to the typical buffer protocol). Is there more to it than that? (I mean, that alone sounds like something I very much want, but I'm not sure if our project is prepared to migrate to nanobind just for that)

galv · 2022-10-07T23:01:25Z

The nb::tensor interface is kind of confusing to read in the nanobind source code because its mixed in with a lot of implementation details, but this page shows some of what you can do with it: https://github.com/wjakob/nanobind/blob/master/docs/tensor.md

Probably the main useful thing is being able to specify compile time shapes for less confusing code. I know it also always fills in the "strides" field if it happens to be null (allowed by dlpack), which avoids the foot gun of indexing into a NULL strides field. (This is non-trivial to do because something has to own the heap-allocated strides array, and the producer of a dlpack tensor, which provides a destructor, can't do it because it didn't create the strides array itself.)

Meanwhile, the DLManagedTensor interface from dlpack.h is a spartan stable ABI that is best described by reading the one header file that defines it: https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h

wjakob · 2022-11-09T20:29:22Z

Where does this task stand? Halide would very much like to support DLPack in our Python bindings, but nanobind isn't (yet) a good option for us.

What's missing in nanobind, just out of curiosity? @steven-johnson

Skylion007 added the enhancement label Apr 11, 2022

Skylion007 changed the title ~~[ENH]: Add DLPACK castersto PyBind11~~ [ENH]: Add DLPACK casters to pybind11 Apr 11, 2022

Skylion007 changed the title ~~[ENH]: Add DLPACK casters to pybind11~~ [ENH]: Add DLPack casters to pybind11 Apr 11, 2022

Skylion007 added the help wanted label Apr 14, 2022

Skylion007 assigned oleksandr-pavlyk Apr 14, 2022

Skylion007 mentioned this issue Apr 14, 2022

Expand dtype accessors #3868

Merged

galv mentioned this issue Oct 7, 2022

fix: PyCapsule_GetDestructor is allowed to return a nullptr destructor #4221

Merged

rwgk mentioned this issue Feb 10, 2023

FWD pybind11 google/pybind11clif#3858

Closed

oandreeva-nv mentioned this issue Mar 24, 2023

Update dlpack implementation for PbTensor triton-inference-server/python_backend#223

Merged

Idclip mentioned this issue Feb 14, 2024

Port OpenVDB Python bindings from pybind11 to nanobind AcademySoftwareFoundation/openvdb#1753

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ENH]: Add DLPack casters to pybind11 #3858

[ENH]: Add DLPack casters to pybind11 #3858

Skylion007 commented Apr 11, 2022 •

edited

Loading

Skylion007 commented Apr 14, 2022

oleksandr-pavlyk commented Apr 14, 2022

PeterDykas commented May 25, 2022

oleksandr-pavlyk commented May 25, 2022

PeterDykas commented May 25, 2022

Tabrizian commented Jul 6, 2022

galv commented Jul 11, 2022 •

edited

Loading

steven-johnson commented Oct 7, 2022

galv commented Oct 7, 2022

steven-johnson commented Oct 7, 2022

galv commented Oct 7, 2022

wjakob commented Nov 9, 2022 •

edited

Loading

[ENH]: Add DLPack casters to pybind11 #3858

[ENH]: Add DLPack casters to pybind11 #3858

Comments

Skylion007 commented Apr 11, 2022 • edited Loading

Required prerequisites

Problem description

Skylion007 commented Apr 14, 2022

oleksandr-pavlyk commented Apr 14, 2022

PeterDykas commented May 25, 2022

oleksandr-pavlyk commented May 25, 2022

PeterDykas commented May 25, 2022

Tabrizian commented Jul 6, 2022

galv commented Jul 11, 2022 • edited Loading

steven-johnson commented Oct 7, 2022

galv commented Oct 7, 2022

steven-johnson commented Oct 7, 2022

galv commented Oct 7, 2022

wjakob commented Nov 9, 2022 • edited Loading

Skylion007 commented Apr 11, 2022 •

edited

Loading

galv commented Jul 11, 2022 •

edited

Loading

wjakob commented Nov 9, 2022 •

edited

Loading