[RFC] Support DLPACK C Functions for Speed Exchange and Stream Handling

[tqchen]

This PR adds support for three C functions to speedup DLPack exchange. As of now, DLPack exchange relies on python functions such as tensor.__dlpack__().

While they works well for common cases, the general overhead of such exchange is at the level of 0.2-0.3 us for very well optimized version, and can go up to 0.4-1 us for less optimized implementation.

For a function that takes three arguments f(a, b, c), assume we run DLPack exchange for each argument, the general conversion overhead usually gets to around 1us and sometimes to 3us.

While such overhead can be acceptable in many settings, in GPU applications the extra 1-3us overhead can still be significant.

This PR proposes four functions for speed exchange DLPack tensors without going through python interpreter.

• DLPackManagedTensorFromPyObjectNoSync for fast exchange owned tensors into consumer
• DLPackDLTensorFromPyObjectNoSync for fast exchange non-owned tensors into consumer
• DLPackManagedTensorToPyObjectNoSync for fast exchange consumer tensors into producer (return values)
• DLPackCurrentWorkStream for query the current working stream from the stream context

Our preliminary results show that these functions, when incorporated correctly via native extensions such as c/c++, can bring exchange cost to the level of 30ns - 80ns, giving us about one order of maginitude speedup. That means the API overhead of functions like f(a, b, c) will be at 0.2us-0.4us level (including exchange), which is close to what native cpp extension overhead do without exchange.

[tqchen]

RFC #175

[tqchen]

updates to incorporate suggestion by @dalcinl.

[tqchen]

This PR is updated to reflect the all the suggestions in #175

Naming (updated per suggestion from @oleksandr-pavlyk @kkraus14):

• DLPackManagedTensorAllocator
• DLPackManagedTensorToPyObject
• DLPackManagedTensorFromPyObject

Clarified that the dunder should be attached to the class type @gbonik and @seberg

[kkraus14]

Everything LGTM except for the ongoing discussions related to stream / synchronization handling

[tqchen]

Thanks everyone for comments. We have updated the proposal to include a non-owned version while explicitly state clearly the intend is for fast caling for framework library/DSL call use cases

Summary of the current functions

• DLPackManagedTensorFromPyObjectNoSync for fast exchange owned tensors into consumer
• DLPackDLTensorFromPyObjectNoSync for fast exchange non-owned tensors into consumer
• DLPackManagedTensorToPyObjectNoSync for fast exchange consumer tensors into producer (return values)
• DLPackCurrentWorkStream for query the current working stream from the stream context

I think it is getting to a mergable state, but would be good to also get final inputs if any

[yongwww]

I’m in favor of this RF. Adding C APIs for DLPack exchange and stream querying is a very compelling direction, achieving exchange latency down to tens of nanoseconds could remove a significant bottleneck. This RFC is a valuable addition to DLPack. I’m excited to see it get landed and deployed in production!

[tqchen]

incoprorated suggestions from @seberg into the PR. i think we are in good shape, going to merge this in two days if there is no more future comments.

Thanks everyone for suggestions so far

[seberg]

I like this approach in general. I am still a bit unsure about the allocation function and the filling one, but if nobody else chimes in, then so be it...

One question about the stream sync, we should clarify that a bit maybe.

One curiosity: Should there be an entry (or function?) that allows discovering which device types are supported?

---

I think in a sense, I still need to see how this actually looks like for a real-world consumer/producer pair, but in the near future it is maybe still malleable enough.

[seberg]

One small thing is, I am not sure how we will use kind? Let's say this typically get's converted to a Python exception, then we would have to somewhat agree on e.g. MemoryError or so to translate that.

Maybe this will just settle itself, in the end I bet this pretty much only ever raises MemoryError anyway.

(I think this function is important! I know I am slow to settle on something, but I haven't quite settled on loving this approach. But I suspect its true that asking to allocate e.g. the torch tensor here -- requiring the GIL -- and then viewing that is also not really better.)

[tqchen]

yes, likely MemoryError is the most common one likely

[seberg]

I like this approach as such, but as I just mentioned it elsewhere, I want to point out that if this is optional and commonly not support, it means all consumers will need:

if (api->managed_tensor_to_py_object_no_sync == NULL) {
    // do complicated thing
}

which unfortunately means that we can get a small speed improvement if available (not having to allocate the DLManagedTensor, etc.) but we may need to support both paths as a consumer, unfortunately.

I suppose part of the solution here may be that you really want a C++ convenience layer that is easy to vendor...

This isn't a deal breaker! It made me wonder if filling in a DLManagedTensor is much worse -- it would be different from currently by not owning its own allocation.

[seberg]

I forgot about the CPU part, allowing the function pointer to be NULL makes sense. But I am also happy to just implement always NULL return or an always -1 with a "NumPy doens't have streams, what's going on" error.

[seberg]

I am a bit unsure if this is specified as well as it needs to be? I.e. I think NULL would be the default stream.

The question is: Is there any need (or not) an "undefined or no synchronization" return value (such as -1)?
If not, we are all good, but if some producer might need this (for whatever reason), then we need to specify this here.

The alternative is that the producer just has to return the default stream (otherwise the consumer has to guess a stream anyway probably in the kernel use-case!).

[tqchen]

i don;t think there is a need in this particular context. returning the default stream is likely more well defined

[tqchen]

thanks @seberg on entry function to discover device supported, as of now no, but if there is a need, we could add such API. although i am not sure if it is strictly needed

[dalcinl]

In the future, a version v2 may have a DLPackExchangeAPI struct with more/different entries than v1, therefore the struct corresponding to v2 may not have the same layout (and thus type name) than the older v1.

I assume that the intention here is that DLPackExchangeAPI will grow, if ever, to the end; additional, the entries currently defined below will never change (not their relative position, nor the layout of each struct).
I that correct? If not, this will require some nasty and inconvenient fixes.

Long story short: this is a great idea, but the implementation as it is now is not flexible to future changes the layout of the various structures. Just double checking we are all in the same page here.

[tqchen]

This is a good point. I think we can say only the two first two fields remains unchanged, will do a patch to clarify this

[tqchen]

updated to move the exchange api stable part to a new struct DLPackExchangeAPIHeader

[dalcinl]

If things ever change after the first two fields, then the C code users will write to deal with older versions will violate the C language strict aliasing rules and entering the undefined behavior territory.

[dalcinl]

Oh, I see you changed things to use a header.

[tqchen]

@dalcinl if it works for you, would be great if you can explicitly approve, thanks

[tqchen]

thanks everyone for valuable feedbacks so far, planning to merge in 24hours

[tqchen]

Thanks everyone, this is merged