feat(pyo3): consumer + producer DLPackExchangeAPI#30
Open
HaoZeke wants to merge 2 commits into
Open
Conversation
Consumer side: adds `ExchangeAPI`, a Copy wrapper around the producer's function table exposed via `type(obj).__dlpack_c_exchange_api__`, with safe methods for all four callback slots: `current_work_stream`, `managed_tensor_from_pyobject`, `managed_tensor_to_pyobject`, and `with_dltensor_from` (closure-style so the producer-owned shape/stride borrow stays scoped). Producer side: the `PyDLPack` pyclass exposes a static `DLPackExchangeAPI` via a `__dlpack_c_exchange_api__` classattr. Three callbacks are wired up to the existing capsule plumbing: `managed_tensor_from_py_object_no_sync` transfers ownership out by renaming the capsule to `used_dltensor_versioned` (same mechanism as the legacy `__dlpack__` path, so a second consume fails cleanly); `dltensor_from_py_object_no_sync` fills a stack DLTensor borrowing the capsule's storage; `managed_tensor_to_py_object_no_sync` rewraps an owned DLPackTensor as a fresh PyDLPack. `current_work_stream` returns null for all devices -- dlpk does not own an accelerator stream, and the upstream producer (if any) is expected to sync before handing the tensor off. Each C callback is panic-caught so a Rust panic cannot unwind across the extern "C" boundary. The capsule name constant moves to `sys::DLPACK_EXCHANGE_API_CAPSULE_NAME`. Tests: a pure-Rust stub producer roundtrips a 1D f32 tensor through the consumer wrapper, and the PyDLPack side roundtrips both through the fast borrow (`with_dltensor_from`) and the fast consume (`managed_tensor_from_pyobject`), with a second consume asserted to fail. `managed_tensor_to_pyobject` is covered by re-exporting the returned PyDLPack through `__dlpack__()`. Closes metatensor#18.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Consumer side: adds
ExchangeAPI, a Copy wrapper around the producer's function table exposed viatype(obj).__dlpack_c_exchange_api__, with safe methods for all four callback slots:current_work_stream,managed_tensor_from_pyobject,managed_tensor_to_pyobject, andwith_dltensor_from(closure-style so the producer-owned shape/stride borrow stays scoped).Producer side: the
PyDLPackpyclass exposes a staticDLPackExchangeAPIvia a__dlpack_c_exchange_api__classattr. Three callbacks are wired up to the existing capsule plumbing:managed_tensor_from_py_object_no_synctransfers ownership out by renaming the capsule toused_dltensor_versioned(same mechanism as the legacy__dlpack__path, so a second consume fails cleanly);dltensor_from_py_object_no_syncfills a stack DLTensor borrowing the capsule's storage;managed_tensor_to_py_object_no_syncrewraps an owned DLPackTensor as a fresh PyDLPack.current_work_streamreturns null for all devices -- dlpk does not own an accelerator stream, and the upstream producer (if any) is expected to sync before handing the tensor off. Each C callback is panic-caught so a Rust panic cannot unwind across the extern "C" boundary.The capsule name constant moves to
sys::DLPACK_EXCHANGE_API_CAPSULE_NAME.Tests: a pure-Rust stub producer roundtrips a 1D f32 tensor through the consumer wrapper, and the PyDLPack side roundtrips both through the fast borrow (
with_dltensor_from) and the fast consume (managed_tensor_from_pyobject), with a second consume asserted to fail.managed_tensor_to_pyobjectis covered by re-exporting the returned PyDLPack through__dlpack__().Closes #18.