[RFC] C Functions for Speed Exchange and Stream Handling

This is a proposal to bring C standard functions pointer APIs to exchange DLPack that can potentially bring 3x-5x speedup to previous best effort cases, and will enable more effective use in GPU settings

Overview

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 0.7us - 3us.

While such overhead can be acceptable in many settings, in GPU applications the extra 1-3us overhead can still be significant. For a kernel that takes 2us to finish, 0.7 us means 30% additional overhead in execution

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

Proposed Functions

/*!
 * \brief A generic C-style allocator that exposes allocation of a Tensor/Array.
 *
 * This information can then be used to set allocators of a callee to run allocations.
 * This information can then be used to set the callee's allocator to perform allocations.
 * This function can be exposed by the framework through the DLPackExchangeAPI.
 *
 * This particular function does not assume a Python environment; as a result,
 * the error handling mechanism is different from Python-related functions.
 *
 * \param prototype The prototype DLTensor to offer details about the device and shape.
 *                  Other field information will be ignored during allocation.
 * \param out The output DLManagedTensorVersioned.
 * \param error_ctx The context to set the error.
 * \param SetError The function to set the error.
 * \return 0 on success, -1 on failure.
 *         The callee should call SetError(error_ctx, kind, message) to set the error kind and message.
 * \note Error propagation via SetError.
 *
 * \sa DLPackExchangeAPI
 */
typedef int (*DLPackManagedTensorAllocator)(                                //
  DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx,     //
  void (*SetError)(void* error_ctx, const char* kind, const char* message)  //
);

/*!
 * \brief Exports a PyObject* Tensor/NDArray to a DLManagedTensorVersioned.
 *
 * This function is a C-style function pointer to quickly convert a PyObject* Tensor/NDArray
 * to a DLManagedTensorVersioned without going through the Python interpreter.
 *
 * This function does not perform any stream synchronization. The consumer should query
 * DLPackCurrentWorkStream to get the current work stream and launch kernels on it.
 *
 * This function is exposed by the framework through the DLPackExchangeAPI.
 *
 * This information can then be picked up by importers and libraries to perform a fast conversion.
 * This function should not throw any exceptions; if it fails, it should return -1 and
 * set the error message via PyErr_SetXXX.
 *
 * \param py_object The Python object to convert; this should be PyObject*.
 *                  We use void* to avoid dependency on Python.h.
 *
 * \param out The output DLManagedTensorVersioned.
 * \return 0 on success, -1 on failure. PyError should be set if -1 is returned.
 * \note We use void* to avoid dependency on Python.h, so this specific type is
 *       not dependent on Python.h and can be copied to dlpack.h.
 *
 * \sa DLPackExchangeAPI, DLPackCurrentWorkStream
 */
typedef int (*DLPackManagedTensorFromPyObjectNoSync)(                 //
  void* py_object,                                                    //
  DLManagedTensorVersioned** out                                      //
);

/*!
 * \brief Imports a DLManagedTensorVersioned to a PyObject* Tensor/NDArray.
 *
 * This function is a C-style function pointer to quickly convert a DLManagedTensorVersioned
 * to a PyObject* without going through the Python Interpreter.
 *
 * This function does not perform any stream synchronization.
 *
 * This function is exposed by the framework through the DLPackExchangeAPI.
 *
 * \param tensor The DLManagedTensorVersioned to convert.
 * \param out_py_object The output Python object.
 * \return 0 on success, -1 on failure. PyError should be set if -1 is returned.
 * \note We use void* to avoid dependency on Python.h, so this specific type is
 *       not dependent on Python.h and can be copied to dlpack.h.
 *
 * \sa DLPackExchangeAPI
 */
typedef int (*DLPackManagedTensorToPyObjectNoSync)(                     //
  DLManagedTensorVersioned* tensor, void** out_py_object                //
);


/*!
 * \brief Exports a PyObject* Tensor/NDArray to a DLTensor whose space is pre-allocated on stack.
 *
 * This function is a C-style function pointer to quickly convert a PyObject* Tensor/NDArray
 * to a DLTensor whose space is pre-allocated on stack without going through the Python interpreter.
 *
 * This is an non-owning conversion, the producer still owns the memory of data, strides, shape.
 * The liveness of DLTensor is only guaranteed until the consumer returns control to the caller.
 *
 * In the context of this function, we expect the producer to allocated space for data, strides and shape.
 *
 * This function does not perform any stream synchronization. The consumer should query
 * DLPackCurrentWorkStream to get the current work stream and launch kernels on it.
 *
 * This function is useful when the consumer do not need to retain the tensor memory.
 * It generally can provide about 2x faster conversion than DLPackManagedTensorFromPyObjectNoSync.
 *
 * For cases where consumer may needs to reorganize the tensor memory via temporary managed copy,
 * DLPackManagedTensorFromPyObjectNoSync should be used.
 *
 * This function is exposed by the framework through the DLPackExchangeAPI.
 *
 * This information can then be picked up by importers and libraries to perform a fast conversion.
 * This function should not throw any exceptions; if it fails, it should return -1 and
 * set the error message via PyErr_SetXXX.
 *
 * \param py_object The Python object to convert; this should be PyObject*.
 *                  We use void* to avoid dependency on Python.h.
 *
 * \param out The output DLTensor, whose space is pre-allocated on stack.
 * \return 0 on success, -1 on failure. PyError should be set if -1 is returned.
 * \note We use void* to avoid dependency on Python.h, so this specific type is
 *       not dependent on Python.h and can be copied to dlpack.h.
 *
 * \sa DLPackExchangeAPI, DLPackCurrentWorkStream
 */
typedef int (*DLPackDLTensorFromPyObjectNoSync)(                      //
  void* py_object,                                                    //
  DLTensor* out                                                       //
);

/*!
 * \brief Obtain the current work stream of a device.
 *
 * This function is a C-style function pointer to obtain the current work stream
 * of a device for frameworks that rely on a context manager to manage the stream.
 * For example, it should map to torch.cuda.current_stream in PyTorch.
 *
 * This function can be set to NULL if the framework does not rely on a context manager
 * to manage the stream. However, we encourage frameworks to provide this function
 * if possible.
 *
 * As if this field is not set, likely consumer cannot safely do stream based
 * exchange based on the
 *
 * \param device_type The device type.
 * \param device_id The device id.
 * \param out_current_stream The output current work stream.
 * \return 0 on success, -1 on failure.
 *
 * \sa DLPackExchangeAPI
 */
typedef int (*DLPackCurrentWorkStream)(                         //
  DLDeviceType device_type,                                     //
  int32_t device_id,                                            //
  void** out_current_stream                                     //
);

/*!
 * \brief Framework-specific function pointers table for DLPack exchange.
 *
 * Guidelines for leveraging DLPackExchangeAPI:
 *
 * There are generally two kinds of consumer needs for DLPack exchange:
 * - N0: library support, where consumer.kernel(x, y, z) would like to run a kernel
 *       with the data from x, y, z. The consumer is also expected to run the kernel with the same
 *       stream context as the producer. For example, when x, y, z is torch.Tensor,
 *       consumer should query exchange_api->current_work_stream to get the
 *       current stream and launch the kernel with the same stream.
 *       This setup is necessary for no synchronization in kernel launch and maximum compatibility
 *       with CUDA graph capture in the producer.
 *       This is the desirable behavior for library extension support for frameworks like PyTorch.
 * - N1: data ingestion and retention
 *
 * Note that obj.__dlpack__() API should provide useful ways for N1.
 * The primary focus of the current DLPackExchangeAPI is to enable faster exchange N0
 * with the support of the function pointer current_work_stream.
 *
 * Array/Tensor libraries should statically create and initialize this structure
 * then return a pointer to DLPackExchangeAPI as an int value in Tensor/Array.
 * The DLPackExchangeAPI* should stay alive throughout the lifetime of the process.
 *
 * One simple way to do so is to create a static instance of DLPackExchangeAPI
 * within the framework and return a pointer to it. The following code
 * shows an example to do so in C++. It should also be reasonably easy
 * to do so in other languages.
 *
 * \code
 * struct MyDLPackExchangeAPI : public DLPackExchangeAPI {
 *   MyDLPackExchangeAPI() {
 *     version.major = DLPACK_MAJOR_VERSION;
 *     version.minor = DLPACK_MINOR_VERSION;
 *     managed_tensor_allocator = MyDLPackManagedTensorAllocator;
 *     managed_tensor_from_py_object_no_sync = MyDLPackManagedTensorFromPyObjectNoSync;
 *     managed_tensor_to_py_object_no_sync = MyDLPackManagedTensorToPyObjectNoSync;
 *     dltensor_from_py_object_no_sync = MyDLPackDLTensorFromPyObjectNoSync;
 *     current_work_stream = MyDLPackCurrentWorkStream;
 *     prev_version_api = nullptr;
 *  }
 *
 *  static const DLPackExchangeAPI* Global() {
 *     static MyDLPackExchangeAPI inst;
 *     return &inst;
 *  }
 * };
 * \endcode
 *
 * Each framework should attach a dunder `__c_dlpack_exchange_api__` integer
 * to point to the DLPackExchangeAPI* pointer.
 *
 * Importantly, the attribute should be attached to the class of the Tensor, not the instance.
 *
 * mypackage.Tensor.__c_dlpack_exchange_api__ = MyPackageDLPackExchangeAPI
 *
 * or equivalently:
 *
 * type(tensor_obj).__c_dlpack_exchange_api__ = MyPackageDLPackExchangeAPI
 */
struct DLPackExchangeAPI {
  /*!
   * \brief The current DLPack version.
   */
  DLPackVersion version;
  /*!
   * \brief Optional pointer to an older DLPackExchangeAPI in the chain.
   *
   * It should be set to NULL if the framework does not support older versions.
   *
   * \sa DLPackExchangeAPI
   */
  struct DLPackExchangeAPI* prev_version_api;
  /*!
   * \brief Framework-specific function pointer for DLPackManagedTensorAllocator
   * \sa DLPackManagedTensorAllocator
   */
  DLPackManagedTensorAllocator managed_tensor_allocator;
  /*!
   * \brief Framework-specific function pointer for DLPackManagedTensorFromPyObject
   * \sa DLPackManagedTensorFromPyObject
   */
  DLPackManagedTensorFromPyObjectNoSync managed_tensor_from_py_object_no_sync;
  /*!
   * \brief Framework-specific function pointer for DLPackManagedTensorToPyObject
   * \sa DLPackManagedTensorToPyObject
   */
  DLPackManagedTensorToPyObjectNoSync managed_tensor_to_py_object_no_sync;
  /*!
   * \brief Framework-specific function pointer for DLPackDLTensorFromPyObject
   * \sa DLPackDLTensorFromPyObjectNoSync
   */
  DLPackDLTensorFromPyObjectNoSync dltensor_from_py_object_no_sync;
  /*!
   * \brief Framework-specific function pointer for DLPackCurrentWorkStream
   *
   * This function can be set to NULL if the framework does not rely on context manager to manage the stream.
   *
   * \sa DLPackCurrentWorkStream
   */
  DLPackCurrentWorkStream current_work_stream;
};

Discussions

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.

Relation with current DLPack API

The proposed speed exchange function can be used in conjunction with the current DLPack exchange, to gracefully handle fallback cases. To ensure exchange happens in very efficient way for common cases that need zero copy, likely we don't expect exchange to happen for all cases. We can still fallback to the normal dlpack exchange that offers more robust and still efficient exchanges

Reference pytorch implementation (in torch extension)

#include <ATen/DLConvertor.h>
#include <ATen/Functions.h>
#include <c10/cuda/CUDAStream.h>

int TorchDLPackFromPyObject(void* py_obj, DLManagedTensorVersioned** out) {
  try {
    py::handle handle(static_cast<PyObject*>(py_obj));
    at::Tensor tensor = handle.cast<at::Tensor>();
    *out = at::toDLPackImpl<DLManagedTensorVersioned>(tensor);
    return 0;
  } catch (const std::exception& e) {
    PyErr_SetString(PyExc_RuntimeError, e.what());
    return -1;
  }
}

int TorchDLPackToPyObject(DLManagedTensorVersioned* src, void** py_obj_out) {
  try {
    at::Tensor tensor = at::fromDLPackImpl<DLManagedTensorVersioned>(src, nullptr);
    *py_obj_out = THPVariable_Wrap(tensor);
    return 0;
  } catch (const std::exception& e) {
    PyErr_SetString(PyExc_RuntimeError, e.what());
    return -1;
  }
}

int TorchDLPackTensorAllocator(
    DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx,
    void (*SetError)(void* error_ctx, const char* kind, const char* message)
) {
  try {
    at::IntArrayRef shape(prototype->shape, prototype->shape + prototype->ndim);
    at::TensorOptions options = at::TensorOptions()
      .dtype(at::toScalarType(prototype->dtype))
      .device(at::getATenDevice(prototype->device.device_type, prototype->device.device_id));
    at::Tensor tensor = at::empty(shape, options);
    *out = at::toDLPackImpl<DLManagedTensorVersioned>(tensor);
    return 0;
  } catch (const std::exception& e) {
    SetError(error_ctx, "TorchDLPackTensorAllocator", e.what());
    return -1;
  }
}


int TorchCurrentWorkStream(DLDeviceType device_type, int32_t device_id, void **out_stream) {
  try {
    if (device_type != kDLCPU) {
       *out_stream = at::cuda::getCurrentCUDAStream(device_id).stream();
    }
    return 0;
  } catch (const std::exception& e) {
    PyErr_SetString(PyExc_RuntimeError, e.what());
    return -1;
  }
}

struct TorchDLPackExchangeAPI : public DLPackExchangeAPI {
   TorchDLPackExchangeAPI() {
       version.major = DLPACK_MAJOR_VERSION;
       version.minor = DLPACK_MINOR_VERSION;
       prev_version_api = nullptr;
       managed_tensor_allocator = TorchDLPackTensorAllocator;
       managed_tensor_from_py_object_no_sync = TorchDLPackFromPyObject;
       managed_tensor_to_py_object_no_sync = TorchDLPackToPyObject;
       dltensor_from_py_object_no_sync = TorchDLTensorFromPyObject;
       current_work_stream = TorchCurrentWorkStream;
   }
   
   const DLPackExchangeAPI* Global() {
       static MyDLPackExchangeAPI inst;
       return &inst;
   }
};


// The following code should be part of the python binding, to expose these function pointers
// python binding, likely as update to torch._C.
int64_t TorchDLPackExchangeAPIPtr() {
    return reinterpret_cast<int64_t>(TorchDLPackExchangeAPI::Global());
}
// ...

PR #174

Some benchmark on torch on three argument function f(x, y, z), benchmark the overhead of API by passing in arguments, doing conversion at cpp extension level, and get into C and return(nop), x, y, z are torch.Tensors with one element, [stream] means we set a custom stream

Torch 2.8 before https://github.com/pytorch/pytorch/pull/162111
--------------------------------------------------------
torch[cpu]                           3.076291084289551e-06 sec/call
torch[cuda]                          3.134942054748535e-06 sec/call
torch[cuda][stream]                  3.104686737060547e-06 sec/call

Torch main(after the performance fix in https://github.com/pytorch/pytorch/pull/162111)
--------------------------------------------------------
torch[cpu]                           2.0400047302246094e-06 sec/call
torch[cuda]                          2.0430326461791994e-06 sec/call
torch[cuda][stream]                  2.0476818084716796e-06 sec/call

--------------------------------------------------------
With the proposed change
--------------------------------------------------------
torch[cpu]                           3.7751197814941404e-07 sec/call
torch[cuda]                          3.788471221923828e-07 sec/call
torch[cuda][stream]                  3.9639472961425783e-07 sec/call

Wouldn't it be better to group all the C function pointers in a container, let say __c_dlpack_api__, rather than using individual attributes for each function? Such container could be a dict or, if the dict lookup cost is prohibitive, a plain tuple. There are to reasons for proposing this. The fist one is avoiding dunder attributes being added over and over again when new functionality is introduced. The second one is about encapsulation: I would argue that the allocator/to/from functions may be related and share common bits of implementation (or even internal runtime state), and perhaps it would have sense to set/override them as a group rather than individually.

@dalcinl

Thanks @dalcinl. I think it is certainly possible to go with one grouped field, There are several ways to go with it

A0: Use separate three fields
A1: Frameworks returning a pointer to a function table , say DLPackExchangeAPI, that is statically allocated and return a global pointer to DLPackExchangeAPI. Use a C Function table and expose through __c_dlpack_exchange_api__ as PyLong_FromVoidPtr. Consumer can check DLPackVersion for compact if needed. The struct will be alive throughout the entire lifetime of process so can be used as a jumptable

struct DLPackExchangeAPI {
    DLPackVersion version;
    DLPackTensorAllocator tensor_allocator;
    DLPackFromPyObject dlpack_from_py_object;
    DLPackToPyObject dlpack_to_py_object;
};

// in framework: c++ example, cython is also possible

struct MyDLPackExchangeAPI : public DLPackExchangeAPI {
   MyDLPackExchangeAPI() {
       tensor_allocator = MyDLPackTensorAllocator;
       dlpack_from_py_object = MyDLPackFromPyObject;
       dlpack_to_py_object = MyDLPackToPyObject
   }
   
   const DLPackExchangeAPI* Global() {
       static MyDLPackExchangeAPI inst;
       return &inst;
   }
};

A1: Use a dict
A2: Use a tuple

Would be good to see what people prefer

@dalcinl

I now agree that A1 perhaps is better, updated proposal and PR #174 to reflect that. thank you @dalcinl !

use a C Function table and expose through c_dlpack_exchange_api as PyLong_FromVoidPtr

Let's please use a custom object for this? And with custom object, I mean a capsule. I really don't like this passing raw pointers around for 1ns of speed difference.

I know that it tends to confuse everyone, but I would also make it type(obj).__dlpack_slots__. These functions are unbound methods and unbound methods belong logically to the type (that also makes it hard to attach them to something else incorrectly).
(There is one downside that type lookups used to rarely take a fast path for failure, but that might even be fixed in Python now should have a look, because this was a larger annoyance also in NumPy.)

As used above, I would call it __dlpack_slots__ or so because that is the correct term in Python C-API nomenclature.

Some thoughts on the actual slots:

NIT: The version was a single struct already, let's use it. I am wondering if this is also the max version supported by the export, I guess that can be assumed.
I think you really need to add feature flags into the slots here, even if unused right now. The version isn't enough/useful, IMO. First, it might change without the slots changing and second that still wouldn't allow us to support multiple versions at once really.
- We can remove slots, although only by making them set an error.
So using the result stream as an output pointer makes sense (I think the current handling in DLPack is fundamentally problematic and the CAI design would have been correct). This does need clarifying of the life-time, though.
- I wonder if it isn't better to just integrate this information into the tensor itself. If this is only needed to be filled in optionally, then I think I like request flags to be passed in.
- In general, I think it's a good idea to pass request flags. Learn from the buffer protocol which has this since forever! This could/should mirror flags on the struct, so I can request a stream (or future thing) only after checking that the request will actually be honored. (Unknown requests are ignored, it is up to the exporter to get it right.)
- The stream lifetime is not specified here. For "kernel" style calls this doesn't matter, but for "permanent export" style use -- which the original DLPack was designed for IMO -- it is important. I think my opinion is that if we have any lifetime guarantees, it seems even more natural to attach the stream to the managed tensor.
  (My gut feeling is that lifetime guarantees are probably OKish in practice -- in the unlikely case that is really hard, I suspect one might be able to get away with default stream.)
You can't pass a void **streamptr unless you know the device type. So you need a way to query the device type independently. That is unless you attach the stream to the Tensor directly.
dlpack_to_pyobject is indeed a useful addition!
I am a bit unsure about DLPackTensorAllocator, need to think about it more.
- I guess I first wonder whether a lower level allocator API might not be just as well. But of course then you do need to deal with deallocation as well.
- Then I wonder if it isn't just as well to create an array and then fetch the DLPack from it, that will mean very similar code paths for out=array and out=None.
- I think partly, I wonder how important the part is that this is available for temporary allocations or result allocations where the result is only known later on and we thus don't want to back to the Python layer to deal with this.
For the allocator the "kind" is a bit unclear to me. From a Python perspective, I would like this to have fairly well defined values to map to Python exceptions? Although, in practice there may just be two interesting ones: MemoryError (out of memory) and ValueError in practice.

(Unrelated, although I am slightly curious about the Pytorch speeds. It seemed to me that torch spends most of it's time inside some dispatching logic, I am sure this will help a little, but will it help, since I would think this dispatching happens much the same in the C++ layer? I mean, this should safe around 100ns per argument, although much more seems mainly due to forcing C rather than the other being intrinsically slower here.)

Let's please use a custom object for this? And with custom object, I mean a capsule.

I have a mayor concern about using capsules: they are way too hard to manipulate in pure Python code. Exposing function pointers as integers make it very easy to use ctypes or CFFI to wrap these pointers and use them in pure Python code. If capsules are used, then you have to deal with extracting the pointer out of the capsule. While it is doable, there is an extra layer of complexity on understanding and accessing the Python C/API and dealing with reference counting issues.
It took me a lot of work to get all the gory details right in mpi4py for testing purpose, and there are still a few obscure things related to Python object reference ownership that work simply because of ctypes magic.
https://github.com/mpi4py/mpi4py/blob/master/test/dlpackimpl.py#L287

All that being said, I'm not opposed to using capsules. Just wanted to point out that it is a bit annoying that what you would get from __dlpack_slots__ is hard to manipulate form pure Python code. While I understand that the primary motivation for DLPack is to import/export buffers to a foreign language (C/C++), it would be good to easily also support the native language (Python) with tools like ctypes or CFFI wrapping raw pointers. CPython is going that way with the introduction of the __buffer__ dunder method.

they are way too hard to manipulate in pure Python code

Sure, I just prefer the extra type safety when it seems easy/fast enough. But either way it wouldn't be the end of the world here to use integers.

But note that the problem with __dlpack__s capsule is really a DLPack problem. The issue is that it wanted (for historic reasons) an owning capsule which additionally adds the annoyance of both an owning and non-owning capsule adding layers of deallocation logic that is almost as annoying in C as it is in Python.

This API will always return a static pointer, which is even specified already. So all of that ownership problem goes away. I.e. there is never a destructor, context, .... You do need PyCapsule_New and PyCapsule_GetPointer via ctypes of course, but nothing else.

CPython is going that way with the introduction of the buffer dunder method.

Well, __buffer__ returns a memoryview, so it is fully type safe. (Although, I have a bit doubt that the choice of __releasebuffer_ and forcing a memoryview return was good; I wish I had paid attention at the time.)

@rgommers

Thanks for the discussions. Some thoughts here

I think you really need to add feature flags into the slots here
One thing that ideally we want to get at the C protocol level is simplicity, so i think version is perhaps the right minimum info that captures everything. We can always fallback to the python __dlpack__ for maximum compatibility path

You can't pass a void **streamptr unless you know the device type.
we require the device type to be the device type and index tied to the tensor, so it won't be an issue

DLPackTensorAllocator
The main motivation is to keep allocation simple while without introducing a buffer protocol(for now)

For the allocator the "kind" is a bit unclear to me.
I agree, likely they should be mostly mapped to "ValueError" and "MemoryError", just want to leave things open for future ones while not having to change the ABI.

Choice of Type for Dunder Attribute

I think the main discussion seems to boils down to using raw int or PyCapsule

B0: Use raw int

This is simple dunder and things are static

B1: Use PyCapsule

Looking at precedence from numpy, seems numpy exposesa PyCapsule named _ARRAY_API for its C function table. The name __slots__ may confuse with the python slots mechanism as the layout of struct may not strictly follow python object slots(although they serve similar purposes). Given numpy calls the function table "api", i think it maybe ok following this existing practices.

How about we say that we require __dlpack_exchange_api__ to be a PyCapsule that stores static pointer(without destructor), and we can get it via

B1a: PyCapsule_GetPointer(__dlpack_exchange_api__, NULL), this is what numpy used
B1b: PyCapsule_GetPointer(__dlpack_exchange_api__, "dlpack_exchange_api"), add name to it

One possible argument for it is perhaps from typing pov PyCapsule type indicate external opaque pointer, and it follows a precedence(numpy practice)

Discussions

Personally I think we can go with B1b, B1a, or B0. Following numpy's practice perhaps B1a or B1b can be used, it does adds one level of indirection, so if the data structure is livefetched per call there can be overhead in the ctypes layer if things are implemented in this way. Love to get takes from others as well.

cc @rgommers @dalcinl @leofang

+1 for using __dlpack_exchange_api__ name over __dlpack_slots__.

I would suggest renaming DLPackTensorAllocator to DLPackManagedTensorCreate or DLPackManagedTensorNew, to better reflect that it creates new copy of DLManagedTensorVersioned struct.

Even if the name is kept as is, this point about it creating a new copy of struct must be emphasized in the doc-string.

Ditto for DLPackToPyObject and DLPackFromPyObject.

Since the purpose of dlpack_from_py_object is to enable faster kernel launches on arbitrary types representing CUDA device-accessible allocations, we either need standardize layout behind opaque pointer env_stream, or this argument never sees adoption.

Happy to open discussion in another issue related to stream ownership semantics / lifetime management if appropriate, let me know.

After a lot of discussion and consideration from the Apache Arrow community in building the Arrow C Device Data Interface, they decided on exposing an event pointer to use for stream / synchronization management: https://arrow.apache.org/docs/format/CDeviceDataInterface.html#c.ArrowDeviceArray.sync_event and that the lifetime of that event is to be controlled as part of the release function.

If I remember correctly, the reason for the __dlpack__ python protocol using a consumer provided stream instead of a producer provided stream was because there wasn't confidence in frameworks being able to expose stream lifetime management / ownership semantics through the dlpack deleter function. I imagine we run into the same problem here with the current proposal?

What are the lifetime guarantees of optional_out_env_stream? Who is responsible for destroying that stream?

@kkraus14

Thanks @kkraus14 for the comment. You are right that thinking about stream management beyond the local scope is still hard.

In this case, we are thinking about the common pattern that major deep learning framework(e.g. PyTorch) and we are limiting the scope here in the same way as __dlpack__ protocol

stream = torch.Stream()

with torch.stream(stream):
     # running on the stream 
     compute(x, y, z)
# running on default context

So reflecting that status, stream ownership is not done by the DLPack itself, but instead managed by the "environment framework". The particular exchange optional_out_env_stream guarantees that if we do dlpack exchange for compute, the following computation should be carried out using the stream. However, the stream is not attached to the Tenaor/Array itself. The ownership of stream still belong to torch.

@oleksandr-pavlyk

@oleksandr-pavlyk, good pt about naming, perhaps we can rename to

DLPackManagedTensorAllocator
DLPackManagedTensorToPyObject
DLPackManagedTensorFromPyObject

I am not too sure about what you mean by layout of env_stream, but in this case, it is supposed to be the native raw stream, aka cudaStream_t in the case of cuda.

* This function can be set to NULL if the framework does not rely on a context manager to manage the stream.

This seems very Python specific where I think we should reword this to indicate that the stream lifetime is expected to be managed outside of DLPack as opposed to specifically talking about a context manager.

* \param optional_out_last_active_stream Outputs the current stream the tensor is synced to.
*   It can be NULL, in which case the stream will not be queried.
*   optional_out_last_active_stream should point to cudaStream_t in the case of CUDA.
*   Note that for frameworks that use a stream context manager, optional_out_last_active_stream
*   can be the stream that the context manager was most recently active on.
*   The stream is owned by the producer, and the consumer cannot retain it.
*   Instead, the consumer can record an event or add wait dependencies to it.
*   It is the responsibility of the consumer to synchronize with the stream if necessary.
*   The producer may output `reinterpret_cast<void*>(-1)` to indicate that the last active stream
*   is not available; in such a case, a device sync is needed to ensure data is ready.

Using a stream here still feels like the wrong approach to me versus having an optional Event member as part of the DLManagedTensorVersioned struct with lifetime / ownership semantics handled by the existing deleter function. I feel there's a number of shortcomings of the current approach in passing a stream:

Strong possibility of a dangling pointer
Passing -1 for when we can't guarantee Stream lifetime forces synchronizing the entire device which is extremely inefficient
No way to pass per-thread default streams through safely unless we guarantee not to use different threads
If producer wants to continue using the Stream being passed for additional work, consumer may get false dependencies or accidentally block unrelated producer work while trying to record the event themselves

I would propose serving both N0 and N1 from a single API (needs wordsmithing but hopefully gets the idea across 😅):

/*!
 * \brief Exports a PyObject* Tensor/NDArray to a DLManagedTensorVersioned and optionally obtains the current work stream
 *
 * This function is a C-style function pointer to quickly convert a PyObject* Tensor/NDArray
 * to a DLManagedTensorVersioned and optionally obtain the current work stream without going
 * through the Python interpreter.
 *
 * This function is exposed by the framework through the DLPackExchangeAPI.
 *
 * This information can then be picked up by importers and libraries to perform a fast conversion.
 * This function should not throw any exceptions; if it fails, it should return -1 and
 * set the error message via PyErr_SetXXX.
 *
 * \param py_object The Python object to convert; this should be PyObject*.
 *                  We use void* to avoid dependency on Python.h.
 *
 * \param out The output DLManagedTensorVersioned.
 *
 * \param optional_out_current_stream The output current work stream.
 *   It can be NULL, in which case no current work stream is supplied. In this case, the
 *   Event member of the output DLManagedTensorVersioned needs to be used to
 *   access its data safely.
 *   If it is not NULL, it is guaranteed to be safe to use the stream to work on the data
 *   in the output DLManagedTensorVersioned without needing to wait on its Event
 *   member.
 *   If both this stream and the Event member are NULL, the data is guaranteed to
 *   be safe to operate on from any Stream.
 *   TODO: Do we need an escape hatch for implementation defined behavior?
 *   optional_out_current_stream should point to cudaStream_t in the case of CUDA.
 *
 * \return 0 on success, -1 on failure. PyError should be set if -1 is returned.
 * \note We use void* to avoid dependency on Python.h, so this specific type is
 *       not dependent on Python.h and can be copied to dlpack.h.
 *
 * \sa DLPackExchangeAPI
 */
typedef int (*DLPackTensorAndStreamFromPyObject)(               //
  void* py_object,                                              //
  DLManagedTensorVersioned** out,                               //
  void** optional_out_current_stream                            //
);

TLDR:

Add optional Event member to DLManagedTensorVersioned and Event lifetime is controlled by existing DLManagedTensorVersioned deleter function
Move API back to returning both DLManagedTensorVersioned and an optional stream expected to be used for consumer kernels.
If both optional_out_current_stream and Event member are not NULL. Data is guaranteed to be safe to use in the stream but the Event can be waited on to access the data safely.
If optional_out_current_stream is not NULL and Event member is NULL, data is guaranteed to be safe to use in that stream but not guaranteed to be safe to use in any other stream.
If optional_out_current_stream is NULL and Event member is not NULL, then Event member must be waited on to access the data safely.
If both optional_out_current_stream and Event member are NULL, then data is guaranteed to be safe to use from any Stream / is synchronized.
TODO: Do we need an escape hatch to define implementation defined synchronization behavior?

This allows for passing Streams around nicely for use cases like Graph Capture but still supports the "ingestion" path where a producing library wants to share only data and not anything related to its execution environment.

@kkraus14

Thanks @kkraus14 . I still think we would benefit from two explicit APIs that distinguishes N0 and N1.

The original DLPackManagedTensorFromPyObject perhaps comes at a case that strongly streamlines N0 while it maybe hard to retrofit to N1.
I think we are going to hit the same issue of complexity comes when we try to start with one API for N1 and retrofit for N0 (like in the case of DLPackTensorAndStreamFromPyObject ).

Specifically coming from context of N0:

N0a: First, it is desirable to limit the times of stream query, so optional_out_last_active_stream can be NULL, which means the consumer gives up query stream, instead relies on DLPackCurrentWorkStream query once. Which helps the critical loop
N0b: It is also desirable for consumer to check whether producer is "stream context manager style", which i would argue a very important piece of design details, since that defines how libraries can work with sucn case
N0c: Because stream only guarantees to work for the period of the calling context (specific to N0), there is no need to worry about carry stream lifetime.

In summary, the N0 consumers have desire to minimize amount of event recording, sync decisions made.

Comparison of Proposals: Should the Event Recording Strategy Falls into Producer or Consumer

Now coming back to the two function proposal. I think in general the final achievable result is that that different from your proposal, except that in the case of DLPackTensorAndStreamFromPyObject the overall decision/duty goes into the producer to decide on whether to allocate event, manage lifecycles and so on.

While in the DLPackManagedTensorFromPyObject + DLPackCurrentWorkStream case, the decision falls onto the consumer:

Consumer can always choose to record an Event on optional_out_last_active_stream, and attach that during ingestion, that is effectively equivalent to the event recording process in N1, except that
It is important that the decision of event recording falls onto the consumer, not the producer, since for the same producer (e.g. torch), depending on the consumer needs N0 or N1:
- N0 consumers DO NOT want to record events
- N1 consumers may wat to record events

Finally specific to note some of the concerns listed. "Shortcomings of the current approach in passing a stream":

Strong possibility of a dangling pointer:

We are limiting the possible scope of stream usage, specifically for N0, it is the current library invocation, for N1, only record during exchange, so there is no issue of dangling pointer as we are not retaining stream.

I think honest take here is that we do not want consumer to retain stream, instead consumer should only use that for the specific kernel invocation following exchange (in case of N0), or recored event/sync to its own stream.

Passing -1 for when we can't guarantee Stream lifetime forces synchronizing the entire device which is extremely inefficient:

That is true, however, i would argue such cases are corner case, as a matetr of fact, i cannot think of an existing framework that it happens. Note that when stream lifetime goes out and event is still provided, we are effectively saying one have to pre-record an event eagerly.

Note that if I am a framework implementer, likely what i would do when stream lifecycle goes out and i still like to have some visible effect to downstream, is to add a sync to some default stream, so future queries can retrun

No way to pass per-thread default streams through safely

I am not sure what is being meant here. I assume passing optional_out_last_active_stream is OK for N1 use-cases to immediately create an event recording. For the case of N0 very likely a separate stream mechanism is used besides the thread default stream

If producer wants to continue using the Stream being passed for additional work, consumer may get false dependencies or accidentally block unrelated producer work while trying to record the event themselves

Consumer should record the event right away at the same thread of the producer after the query in the case of N1, so likely won't block unrelated producer work when the call returns to the producer.

I also want to followup on a possible details that needs to be call out. For stream context manager based framework like PyTorch. It is generally not possible to guarantee the data availability under the current working stream abitrary. Consider the code below:

s0 = torch.cuda.Stream()
s1 = torch.cuda.Stream()
a = torch.empty(...)
with torch.stream(s0):
     b = a + 1

with torch.stream(s1):
    # read b
    # note that we do not guarantee that b is readily available.
    # Additionally, even if we replace mylib.kernel by torch.kernel, it is still a bug
    b_imported = mylib.from_dlpack(b)
    mylib.kernel(b_imported)

In this case, the current working stream in the s1 context is s1. However, the above code do not guarantee that b_imported is synced to s1. Note that even if we replace mylib.kernel by a torch.kernel, it still causes a bug. So for frameworks that comes with stream context manager, the implicit expectation is that stream sync is done by the user, so a typical correct code looks like

s0 = torch.cuda.Stream()
s1 = torch.cuda.Stream()
event = torch.cuda.Event()
a = torch.empty(...)
with torch.stream(s0):
     b = a + 1
     event.record()

with torch.stream(s1):
    s1.wait_event(event)
    # This is now a correct pattern
    b_imported = mylib.from_dlpack(b)
    mylib.kernel(b_imported)

So there must be some pre-context to the statement like the follows (which is not always true):

it is guaranteed to be safe to use the stream to work on the data in the output DLManagedTensorVersioned

Instead, usually what is stated instead is that "usually the expectation is that the user will take responsibility to synchronize the data to the current working stream before hand, and for N0 usecases, it is desirable to run compute on the current stream".

Again given stream context and user-managed dependency is something that is common for our key use-cases, i think it is important to call it out, and have a path to optimize for that

I'm going to step away from this issue as I think the arguments are going in circles at this point and I'm not helping to make forward progress. Thanks for all of your hard work and efforts here, they are appreciated!

Following up on this, capturing more discussions. We have updated the proposal to focus on N0 needs. I think the proposal is getting to a ready state, please share your thoughts.

Thanks everyone for discussions so far

sending a note here, we are likely going to merge #174 soon

I'm surprised this would be considered for merging into the standard this year. It seems that there has recently been significant development work, with at least one contributor having to step away from this issue. Shouldn't we see buy-in from stakeholders and at least tentative implementations before proceeding to standardization?

@hpkfft

Thanks @hpkfft for the note. The original contention point is to support both N0 and N1, and as of now the latest update primarily focused on N0 while leaving N1 topics for future.

This is a chicken and egg problem, as we need to settle on at least one version of the standard before bring up the implementations in the downstreams stakeholders. We did discuss the proposal at the array api meeting a few weeks ago and generally the agreement is to go forward. In terms of development, we are already converging on the technical parts that focusing on N0.

Given the significant performance improvement it can bring, and further changes are still compatible through verisoning. I see it is overall postive to roll it out to make dlpack the de-facto choice for the N0 scenarios.

#174 is now merged, we will followup with 1.2 release after #177 get merged

[DLPack] C Functions for DLPack Speed Exchange and Stream Handling (#96)

tagged in https://github.com/dmlc/dlpack/releases/tag/v1.2

[DLPack] C Functions for DLPack Speed Exchange and Stream Handling (#…

[DLPack] C Functions for DLPack Speed Exchange and Stream Handling (p…

[DLPack] C Functions for DLPack Speed Exchange and Stream Handling (#…

Navigation Menu

Repository navigation

[RFC] C Functions for Speed Exchange and Stream Handling #175

Overview

Proposed Functions

Discussions

Relation with current DLPack API

Choice of Type for Dunder Attribute

B0: Use raw int

B1: Use PyCapsule

Discussions

7 remaining items

Comparison of Proposals: Should the Event Recording Strategy Falls into Producer or Consumer

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[RFC] C Functions for Speed Exchange and Stream Handling #175

Description

Overview

Proposed Functions

Discussions

Relation with current DLPack API

Activity

tqchen commented on Sep 12, 2025

tqchen commented on Sep 12, 2025

tqchen commented on Sep 12, 2025

dalcinl commented on Sep 14, 2025

tqchen commented on Sep 14, 2025

tqchen commented on Sep 14, 2025

seberg commented on Sep 16, 2025

dalcinl commented on Sep 16, 2025

seberg commented on Sep 16, 2025

tqchen commented on Sep 16, 2025

Choice of Type for Dunder Attribute

B0: Use raw int

B1: Use PyCapsule

Discussions

oleksandr-pavlyk commented on Sep 17, 2025

kkraus14 commented on Sep 17, 2025

tqchen commented on Sep 17, 2025

tqchen commented on Sep 17, 2025

7 remaining items

kkraus14 commented on Sep 30, 2025

tqchen commented on Sep 30, 2025

Comparison of Proposals: Should the Event Recording Strategy Falls into Producer or Consumer

tqchen commented on Sep 30, 2025

kkraus14 commented on Sep 30, 2025

tqchen commented on Oct 6, 2025

tqchen commented on Oct 9, 2025

hpkfft commented on Oct 9, 2025

tqchen commented on Oct 9, 2025

tqchen commented on Oct 11, 2025

tqchen commented on Oct 11, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions