owned this note
owned this note
Published
Linked with GitHub
Extended buffer protocol PEP draft for device memory support
============================================================
Abstract
========
The Python buffer protocol is widely adopted across the ecosystem to share data between packages such as NumPy and its downstream libraries. Adoption for example happens via Cython's [`typed memoryviews`](https://cython.readthedocs.io/en/stable/src/userguide/memoryviews.html).
However, it was not extended for many years and modern scientific programs now often use accelerators such as GPUs.
In part this need has created many new protocols, such as:
* [Cuda array interface](https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html)
* [DLPack](https://dmlc.github.io/dlpack/latest/)
* [Arrow Data Interface](https://arrow.apache.org/docs/format/CDeviceDataInterface.html)
and these solve most of the issues that occur to varying degree. Unfortunately, none have quite the reach and low integration as the Python buffer protocol.
In this PEP we propose a light-weight extension of the buffer protocol, to allow exposing non CPU buffer.
This extension is designed so that it will allow/simplify addition of new features.
An example of such an extension would be indicating buffer borrowing that can be useful for ownership management and is desired for example in rust.[^and-gpus]
Motivation
==========
Todays data intensive workflows often use CPU and accelerators or cross boundaries between between these.
We believe that the buffer protocol could fill at many of these needs if it were to be extended.
Importantly, we believe doing this inside the buffer protocol will help projects that wish to support a wide variaty of devices to do so with less code duplication.
The buffer protocol is deeply integrated into Python and because of that has a centrality and at least small performance benefits.
As it is also widely adopted, we believe that if it can fill new needs more widely, it's adoption will increase and make the creation of future protocols unnecessary or simpler.
There is always a danger of creating `N + 1` protocols to do a similar task and as such pushing for more larger buffer protocol use.
Right now, no single protocol solves all needs. Even with this extention, the buffer protocol can cover many, but not all, use-cases.
In general, we believe that the possibility of only partitial adoption should not be seen as problematic.
Libraries may implement whichever subset of the buffer-protocol that is useful and easy to implement for them.
This proposal not just allows adoption by more current use-cases, it also unblocks extending the most central buffer exchagne protocol in Python to push new ideas and capabilities.
Rationale
=========
The buffer protocol, with all it's flaws, is widely adopted in the scientific python ecosystem. However, due to it living in Python and the non-obvious nature of how to extend it, new capabilities were never added.
This PEP wishes to address this by:
* Proposing an extension that will simplifier at future additions.
* Enable exporting data that is not directly CPU accessible. This could be in the form of accelerator device memory, distributed, or even compressed data.
A major point in extending the buffer protocol in this direction is that we wish to remain as compatible as possible:
* Any existing code should keep working unmodified (backwards compatibility)
* Ideally, most extension can be used on old Python versions (forward compatibility to the extend possible). This is now and also for future extensions.
Further, we realize that neither Python nor this PEP can or should describe how data exchange on non-CPU memory can work.
Exchanging data on accelerators (or distributed data, ...) is far more complex than for CPU data:
* There are a host of accelerators, all with their own ways or interesting additional information.
* Accelerators often work asynchronously: A function using a dataset may still be working when a new function is already called (similar to threading). This requires additional logic to ensure correct execution order and this logic will be specific to the data.
Thus, the design rationale here is to extend the buffer protocol in backwards compatible way with forward compatibility in mind.
Additionally we leave any device specific information outside Python's purview. Python may facilitate the exchange of such extension, but is unlikely to define them itself.
Specification
=============
New slot for supported flags
----------------------------
While the below device flag specifically accepts ignoring this, we propose adding a new field to the buffer slots to `PyBufferProcs`:
```
struct PyBufferProcs {
/* get buffer and release buffer slots */
int bf_supported_flags;
}
```
A new `PyBUF_REQUIRED_FLAGS` constant (see later) and a way to query these
via:
```
/*
* Check if an object supports the buffer protocol and certain flags.
*
*/
int PyObject_CheckBufferSupports(PyObject *obj, int flags);
```
On future Python versions, `int PyObject_CheckBufferSupports(obj, flags)` can be used to query whether an object is a buffer, and if it is whether it advertises support for all flags (users may need to call `CheckBuffer` additionally as a 0 return may indicate no buffer protocol support).
In practice, some flags are soft request/capability flags. This is currently `PyBUF_INDIRECT`, which indicates support by the consumer but in practice the producer is unlikely to use it.
In practice, Python will define `PyBUF_REQUIRED_FLAGS` to indicate which flags must not be ignored by the producer and `PyObject_GetBuffer` will use `PyObject_CheckBufferSupports(obj, flags & PyBUF_REQUIRED_FLAGS)` and set a `BufferError`.
Currently, it is typical practice to ignore unknown flags and this practice is useful for extension. In the future, Python will enforce flag support for known flags.
**Compatibility**
Currently, supporting only `PyBUF_SIMPLE == 0` is possible, but for types doing nothing `bf_supported_flags == 0` must be the default for technical reasons.
This means that `0` will be translated to support for all currently existing flags and does not change current behavior.
We suggest using `-1` as an indicator that only `PyBUF_SIMPLE` is supported. (One only needs to ensure that `-1` can never have a meaning of "all flags", which may limit the sign bit to be a flag in the future.)
With the above, this extension has the following backwards/foward compatibility:
* Users do not need to check the flags or set the `bf_support_flags` slot to keep existing implementations working. Implementing it would only avoid `BufferError` creation.
* *After* adoption of the new slow, new required flags _can_ be backported (i.e. adopted earlier), with the caveat that the user must check for new flags in `PyObject_CheckBufferSupports` themselves as `PyObject_GetBuffer` would not check for these on older Python versions.
Extended buffer struct and `PyBUF_DEVICE` flag
----------------------------------------------
We propose a new "extended" set of flags with the only current member of the family being `PyBUF_DEVICE`.
The new flag `PyBUF_DEVICE` will be a request flag passed to `PyObject_GetBuffer`.
This flag is *not* a "required" flag, but allows the producer to fill in device information if desired.
If `PyBUF_DEVICE` (or any future extended flag) is passed, the structure passed must have a layout of:
```
struct Py_buffer {
/* Until here identical to previous buffer interface */
int flags;
int ext_flags = 0;
/* Identifier and small scratch space for device indication */
char *device;
void *device_info;
/* Future flags can enable new fields here fields */
}
```
Which currently adds `flags`, to indicate which of the new request flags were used, `ext_flags` as general flag space for the future, and fields to pass device information.
A producer *may* fill in this extended information if the corresponding extended flag is passed. A producer *must not* touch additional fields if the corresponding request flag was not passed.
Thus if `PyBUF_DEVICE` wasn't passed but is required to describe the buffer a `BufferError` must be raised.
(Support must be indicated in `bf_supported_flags`, but this only allows the consumer to avoid calling `PyObject_GetBuffer` if it would reject all CPU buffers anyway.)
If a producer fills in any extended information it *must* set the `flags` to include this information.
That also means that the consumer *must* check the `flags` before using any of the passed fields, even if the producer advertises support.[^device-ping-pong]
In future Python versions `PyObject_GetBuffer()` will zero both `flags` and `ext_flags` to ensure correctness.
**Compatibility**
As producers are free to ignore the extended flags, this extension is fully backwards compatible.
Producers that error on undefined flags may exist, however, we are not aware of any.
One correct observation is that some `PyBuffer_*` functions will only be valid on non-device buffers. However, they cannot be called accidentally, so that this only requires documentation for actual support.
Future Python versions will check the device flag to guard against such use.
Backportability of extensions
-----------------------------
One problem with proposing changes to a protocol such as the buffer protocol is that it will take about 3 years before all typically used Python versions also support it.
This hinders adoption and lowers the incentive to propose changes.
The path of extending the buffer protocol proposed here lowers this issue.
After introducing the proposed changes the `bf_supported_flags` flag will allow backporting new features by careful users even when these features are API incompatible.
Note that the device extension proposed here may also be backported by a careful user as long as they add the necessary checks/preparation proposed for `PyObject_GetBuffer()`.
Device information
------------------
If `PyBUF_DEVICE` is passed, the device information must be filled in, in a well defined way.
Python reserves the `"cpu"` identifier for possible future extension.
Since we reject the idea of Python defining device specific standards, we instead propose that the above mentioned `device` field must be either `NULL` or point to a unique, null terminated utf-8 encoded `char *`.[^unique-char]
If a device is matched, the device specific `device_info` can then be reinterpreted to whatever matches the corresponding specification.
The actual device specification will *not* be specified by Python itself and may be more complex.
As an example, a "CUDA" specification may include the CUDA stream (a concept to allow correct order of operations), a device ordinal to indicate which GPU is used, and additional information about the allocation.
**Choice and list of device identifiers**
Since this proposal is to use a unique name as a device identifier there is a problem of competing naming and authority to use a canonical name.
Python cannot fully control this, but users specifying an extension should open a documentation PR to Python before they adopt a name and Python does reserve the right to reject this.
For example using `SYCL`, `CUDA`, or `HIP` as a name is not acceptable without a clear consensus.
Unless there is a clear authority or consensus names should include a defining PyPI package name to avoid confusion.
Consumers can support multiple definitions, but the producer cannot deprecate theirs, unfortunately.
Notes for device specification
------------------------------
Since an exporter cannot support multiple devices (except via a global config), care should be taken when designing device specific information.
* Include flags or a version to possibly extend the device information in the future.
* You may also wish to reserve additional space to simplify future additions.
Backwards Compatibility
=======================
This PEP is fully backwards compatible. It additionally tries to make the protocol more forward compatible as well.
While Python does not encourage this, the proposed device extension may be adopte in a (mostly) forward compatible way.
The above sections contain brief notes on backwards and forward compatibility.
Reference Implementation
========================
Beyond documentation, this PEP requires small extensions to CPython itself. We propose:
* Define the `PyBUF_DEVICE` flag
* Define `PyBUF_REQUIRED_FLAGS`.
* Defining the new `PyObject_CheckBufferSupports` function.
* Extend the `Py_buffer` struct and allow growing it in the future.
* Changing `PyObject_GetBuffer` to zero out `flags` and `ext_flags` and to check `PyBUF_REQUIRED_FLAGS`.
* Change any API working on buffers (e.g. making contiguous copies) to check that the device is `NULL`.
* Change type creation to fill in `bf_supported_flags` correctly.
Definitions which do not touch `bf_supported_flags` may or may not be backported e.g. to https://github.com/vstinner/pythoncapi_compat.
Additionally, the buffer protocol documentation needs to be extended with these definitions and add notes to all public API functions that cannot work with the new extended flag.
Security Implications
=====================
There are no security implications beyond incorrect implementations.
Rejected Ideas
==============
We are not aware of alternative approaches to extend the buffer protocol itself.
One could add a new slot, but we prefer this adaption for easier future (and current) forward compatibility.
The main alternative would be to invent yet another protocol outside of Python, but as said, we believe that the extension of a widely used protocol is desirable.
Open Discussions
================
How to exchange community defined device types
----------------------------------------------
Python itself should not need to define exact device types as
there are many such devices and they may be complex.
A problem is how to reserve names for new devices and exchange existing definitions.
We propose here to have light review by asking for a PR to the Python documentation as well as discussion on the Python discuss for new ideas.
Unfortunately, there is no way for Python to strictly control this.
Add a new `Py_buffer_extended` struct
-------------------------------------
Rather than extending the actual `Py_buffer` struct
in the future, we could add a `Py_buffer_extended` struct.
We believe that the structure must be extended to at least
include the new flags fields.
In principle, the `device` space itself could be part of
such a `Py_buffer_extended` struct.
We decided that a moderate growth of the structure (two pointer fields), is worth for avoiding two struct definitions.
Including new fields now and future struct growth
-------------------------------------------------
We could include some reserved space already now in the `Py_buffer` struct, which may be nice for future adoption.
However, other than this space coming before the device space there seems little gain in it.
If the `Py_buffer` grows additional fields in the future, users can backport the definition they wish to use these fields before Python includes them.
Add an explicit `PyBUF_EXTENDED` flags
--------------------------------------
We could add an explicit `PyBUF_EXTENDED` request flag to
indicate exactly that the extended flags are available.
This seemed unnecessary to us and we rather extend the main `Py_buffer` struct to always include these flags field.
[^and-gpus]: The concept may also be useful for device data exchange, since knowing that a buffer is only borrowed temporarily can simplify worries about synchronization (where multiple works might use data at the same time).
[^device-ping-pong]: I.e. if a consumer passes `PyObject_GetBuffer(obj, Py_buf_extended &buf, PyBUF_DEVICE);` it must check `(buf->flags & PyBUF_DEVICE)` and if not set, must not access the device specific information and assume a CPU buffer.
[^unique-char]: We choose a string as identifier to allow error messages to include the device information exactly. If applicable, a device specification may wish to use a single unique `char *` pointer to allow.
[^arrow-event]: This is non-trivial and often not a "one size fits all". For example all current GPU protocols differ in their exact exchange of synchronization information, with the safest definition probably the arrow one: https://arrow.apache.org/docs/format/CDeviceDataInterface.html#synchronization.
[^its-complicated] The reason is simply that stream ordering semantics, etc. are complicated. Trying to hash them out would just stall a PEP.