Repository navigation
Allow NumPy data types as arguments #2335
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
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -75,6 +75,58 @@ | |
| * - :py:class:`dpctl.global_mem_cache_type` | ||
| - An :class:`enum.Enum` of global memory cache types for a device. | ||
|
|
||
| .. rubric:: Data types of kernel arguments | ||
|
|
||
| .. _dpctl_data_types: | ||
|
|
||
| Every :py:mod:`dpctl` function or class taking a data type, such as | ||
| :py:meth:`dpctl.SyclQueue.copy`, :py:meth:`dpctl.SyclQueue.fill`, | ||
| :py:class:`dpctl.LocalAccessor`, :py:class:`dpctl.WorkGroupMemory` and | ||
| :py:class:`dpctl.program.SpecializationConstant`, accepts it in any form | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nit: There is no |
||
| :py:class:`numpy.dtype` accepts, e.g. ``"f4"``, ``"float32"``, | ||
| :py:class:`numpy.float32` or a :py:class:`numpy.dtype`. Any data type whose | ||
| values are plain bytes is supported, including structured data types and data | ||
| types registered with :py:mod:`numpy` by other packages, such as those of | ||
| ``ml_dtypes``. Extended precision ``numpy.longdouble`` is not supported. | ||
|
|
||
| A NumPy scalar kernel argument is passed as its bytes, so the kernel determines | ||
| how they are interpreted. | ||
|
|
||
| .. code-block:: python | ||
|
|
||
| import ml_dtypes | ||
|
|
||
| bf16 = ml_dtypes.bfloat16(1.5) # sycl::ext::oneapi::bfloat16 | ||
| c = ml_dtypes.complex32(1 - 2j) # std::complex<sycl::half> | ||
|
|
||
| q.submit(kernel, [x_usm, bf16, c], [n]) | ||
|
|
||
| Scalars of 1, 2, 4, 8 or 16 bytes are always supported. Scalars of other sizes, | ||
| such as a structured scalar matching a C++ ``struct``, require the | ||
| ``raw_kernel_arg`` DPC++ extension. | ||
|
|
||
| Python scalars are converted to a NumPy scalar of a fixed data type, which must | ||
| match the type of the kernel parameter: | ||
|
|
||
| .. list-table:: | ||
| :widths: 20 30 30 | ||
|
|
||
| * - Python type | ||
| - NumPy type | ||
| - C++ type | ||
| * - :py:class:`bool` | ||
| - :py:class:`numpy.bool` | ||
| - ``bool`` | ||
| * - :py:class:`int` | ||
| - :py:class:`numpy.int64` | ||
| - ``std::int64_t`` | ||
| * - :py:class:`float` | ||
| - :py:class:`numpy.float64` | ||
| - ``double`` | ||
| * - :py:class:`complex` | ||
| - :py:class:`numpy.complex128` | ||
| - ``std::complex<double>`` | ||
|
|
||
| .. rubric:: Exceptions | ||
|
|
||
| .. autosummary:: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -81,3 +81,5 @@ profiler | |
| keepdims | ||
| broadcasted | ||
| broadcastable | ||
| bfloat | ||
| dtypes | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| # Data Parallel Control (dpctl) | ||
| # | ||
| # Copyright 2026 Intel Corporation | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| """Data type handling shared by the dpctl functions taking a data type.""" | ||
|
|
||
| import struct | ||
|
|
||
| import numpy as np | ||
|
|
||
| # kinds of NumPy data types whose values are not plain bytes | ||
| _NON_POD_KINDS = frozenset("OSUMmT") | ||
|
|
||
| # extended precision types, which devices do not support | ||
| _EXTENDED_PRECISION_NUMS = frozenset( | ||
| np.dtype(t).num | ||
| for t in (np.longdouble, np.clongdouble) | ||
| if np.finfo(t).nmant > np.finfo(np.float64).nmant | ||
| ) | ||
|
|
||
| # struct formats of the data types whose values are range checked when packed | ||
| _STRUCT_FORMATS = { | ||
| ("b", 1): "?", | ||
| ("i", 1): "b", | ||
| ("i", 2): "h", | ||
| ("i", 4): "i", | ||
| ("i", 8): "q", | ||
| ("u", 1): "B", | ||
| ("u", 2): "H", | ||
| ("u", 4): "I", | ||
| ("u", 8): "Q", | ||
| ("f", 2): "e", | ||
| ("f", 4): "f", | ||
| ("f", 8): "d", | ||
| ("c", 8): "ff", | ||
| ("c", 16): "dd", | ||
| } | ||
|
|
||
|
|
||
| def is_supported_dtype(dt): | ||
| """Returns whether values of :class:`numpy.dtype` ``dt`` can be used on a | ||
| device.""" | ||
| return not ( | ||
| dt.hasobject | ||
| or dt.kind in _NON_POD_KINDS | ||
| or dt.itemsize == 0 | ||
| or dt.num in _EXTENDED_PRECISION_NUMS | ||
| ) | ||
|
|
||
|
|
||
| def as_dtype(dtype): | ||
| """Returns ``numpy.dtype(dtype)``. | ||
|
|
||
| Raises: | ||
| ValueError: if ``dtype`` is not a supported data type. | ||
| """ | ||
| try: | ||
| dt = np.dtype(dtype) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
import dpctl, dpctl.memory as dpm
q = dpctl.SyclQueue()
src, dst = dpm.MemoryUSMShared(16, queue=q), dpm.MemoryUSMShared(16, queue=q)
q.memset(src, 1, 16)
q.memset(dst, 0, 16)
q.copy(dst, src, 2, dtype=None)
print(bytes(dst.copy_to_host()).count(1), "bytes copied") # 16; master: ValueError
q.fill(dst, 1, 2, dtype=">i4")
print(bytes(dst.copy_to_host())[:8].hex()) # 0100000001000000, not big-endianI suggest raising |
||
| except TypeError: | ||
| raise ValueError(f"Unrecognized dtype '{dtype}'") from None | ||
| if not is_supported_dtype(dt): | ||
| raise ValueError(f"Unsupported dtype '{dt}'") | ||
| return dt | ||
|
|
||
|
|
||
| def as_scalar(value): | ||
| """Converts a Python ``bool``, ``int``, ``float`` or ``complex`` to a | ||
| ``numpy.bool``, ``numpy.int64``, ``numpy.float64`` or ``numpy.complex128`` | ||
| respectively. Other values are returned unchanged. | ||
|
|
||
| Raises: | ||
| OverflowError: if an ``int`` does not fit in ``numpy.int64``. | ||
| """ | ||
| if isinstance(value, np.generic): | ||
| return value | ||
| if isinstance(value, bool): | ||
| return np.bool_(value) | ||
| if isinstance(value, int): | ||
| return np.int64(value) | ||
| if isinstance(value, float): | ||
| return np.float64(value) | ||
| if isinstance(value, complex): | ||
| return np.complex128(value) | ||
| return value | ||
|
|
||
|
|
||
| def pack_value(value, dt): | ||
| """Returns the bytes of ``value`` converted to :class:`numpy.dtype` ``dt``. | ||
|
|
||
| Raises: | ||
| ValueError: if ``value`` can not be represented in ``dt``. | ||
| """ | ||
| fmt = _STRUCT_FORMATS.get((dt.kind, dt.itemsize)) | ||
| try: | ||
| if fmt is None: | ||
| return np.asarray(value, dtype=dt).tobytes() | ||
| if dt.kind == "c": | ||
| cval = complex(value) | ||
| return struct.pack("=" + fmt, cval.real, cval.imag) | ||
| return struct.pack("=" + fmt, value) | ||
| except (struct.error, OverflowError, TypeError, ValueError) as e: | ||
| raise ValueError( | ||
| f"Value {value!r} cannot be represented as dtype '{dt}': {e}" | ||
| ) from None | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nit:
DPCTL_BYTES16_Tis described three times: in the NOTE, under Added, and here. This line repeats the NOTE and can go.