dpctl icon indicating copy to clipboard operation
dpctl copied to clipboard

Support `kDLCPU` DLDeviceType for `from_dlpack` and `__dlpack__`

Open ndgrigorian opened this issue 6 months ago • 4 comments

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?

ndgrigorian avatar Aug 02 '24 21:08 ndgrigorian