dpctl
dpctl copied to clipboard
Support `kDLCPU` DLDeviceType for `from_dlpack` and `__dlpack__`
This PR proposes supporting the kDLCPU
DLDeviceType as a targeted device in from_dlpack
and __dlpack__
.
This is accomplished by __dlpack__
by copying the array data to the host, wrapping it in a NumPy array, and then returning a capsule for this NumPy array.
In from_dlpack
, a NumPy array is also returned for devices with __dlpack_device__= (DLDeviceType.kDLCPU, 0)
. This means that from_dlpack(x_np)
for some NumPy array x_np
is effectively a no-op with extra validation steps returning a new Python object wrapping the same memory with zero-copy.
New utility functions numpy_to_dlpack_versioned_capsule
, _numpy_array_interface_from_dl_tensor
, and the class _numpy_array_interface_wrapper
are introduced to enable the new functionality.
This PR also refactors from_dlpack_versioned_capsule
into from_dlpack_capsule
.
- [x] Have you provided a meaningful PR description?
- [ ] Have you added a test, reproducer or referred to an issue with a reproducer?
- [x] Have you tested your changes locally for CPU and GPU devices?
- [x] Have you made sure that new changes do not introduce compiler warnings?
- [ ] Have you checked performance impact of proposed changes?
- [x] If this PR is a work in progress, are you opening the PR as a draft?