Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Expose DLDeviceType enumerator in dpctl.tensor #1793

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions docs/doc_sources/api_reference/dpctl/tensor.constants.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
.. _dpctl_tensor_constants:

Constants
========================

The following constants are defined in :py:mod:`dpctl.tensor`:

.. currentmodule:: dpctl.tensor

.. autodata:: DLDeviceType

.. data:: e

``float``:
IEEE 754 floating-point representation of Euler's constant.

.. data:: inf

``float``:
IEEE 754 floating-point representation of (positive) infinity.

.. data:: nan

``float``:
IEEE 754 floating-point representation of Not a Number (NaN).

.. data:: newaxis

``NoneType``:
Alias for ``None`` which is useful for indexing.

.. data:: pi

``float``:
IEEE 754 floating-point representation of the mathematical constant π.
2 changes: 2 additions & 0 deletions docs/doc_sources/api_reference/dpctl/tensor.rst
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ This module contains:
* :ref:`sorting functions <dpctl_tensor_sorting_functions>`
* :ref:`statistical functions <dpctl_tensor_statistical_functions>`
* :ref:`utility functions <dpctl_tensor_utility_functions>`
* :ref:`constants <dpctl_tensor_constants>`


.. toctree::
Expand All @@ -48,3 +49,4 @@ This module contains:
tensor.sorting_functions
tensor.statistical_functions
tensor.utility_functions
tensor.constants
3 changes: 2 additions & 1 deletion dpctl/tensor/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@
from dpctl.tensor._reshape import reshape
from dpctl.tensor._search_functions import where
from dpctl.tensor._statistical_functions import mean, std, var
from dpctl.tensor._usmarray import usm_ndarray
from dpctl.tensor._usmarray import DLDeviceType, usm_ndarray
from dpctl.tensor._utility_functions import all, any, diff

from ._accumulation import cumulative_logsumexp, cumulative_prod, cumulative_sum
Expand Down Expand Up @@ -376,4 +376,5 @@
"nextafter",
"diff",
"count_nonzero",
"DLDeviceType",
]
36 changes: 36 additions & 0 deletions dpctl/tensor/_usmarray.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,40 @@ include "_slicing.pxi"


class DLDeviceType(IntEnum):
"""
An ``IntEnum`` for the types of DLDevices supported by the DLPack
protocol.
``kDLCPU``:
CPU (host) device
``kDLCUDA``:
CUDA GPU device
``kDLCUDAHost``:
Pinned CUDA CPU memory by cudaMallocHost
``kDLOpenCL``:
OpenCL device
``kDLVulkan``:
Vulkan buffer
``kDLMetal``:
Metal for Apple GPU
``kDLVPI``:
Verilog simulator buffer
``kDLROCM``:
ROCm GPU device
``kDLROCMHost``:
Pinned ROCm CPU memory allocated by hipMallocHost
``kDLExtDev``:
Reserved extension device type used to test new devices
``kDLCUDAManaged``:
CUDA managed/unified memory allocated by cudaMallocManaged
``kDLOneAPI``:
Unified shared memory allocated on a oneAPI non-partitioned device
``kDLWebGPU``:
Device support for WebGPU standard
``kDLHexagon``:
Qualcomm Hexagon DSP
``kDLMAIA``:
Microsoft MAIA device
"""
kDLCPU = c_dlpack.device_CPU
kDLCUDA = c_dlpack.device_CUDA
kDLCUDAHost = c_dlpack.device_CUDAHost
Expand Down Expand Up @@ -1244,6 +1278,8 @@ cdef class usm_ndarray:
The tuple describes the non-partitioned device where the array has been allocated,
or the non-partitioned parent device of the allocation device.

See ``DLDeviceType`` for a list of devices supported by the DLPack protocol.

Raises:
DLPackCreationError:
when the ``device_id`` could not be determined.
Expand Down
Loading