diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cc5c4bf8a..b74561846b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,9 +6,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [dev] - XXX. XX, XXXX +**NOTE**: This release contains breaking changes: + +* The value of the `DPCTL_UNSUPPORTED_KERNEL_ARG` enumerator of `DPCTLKernelArgType` increased by one, as `DPCTL_BYTES16_T` was added before it [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) + ### Added +* `dpctl.SyclQueue.submit` and `dpctl.SyclQueue.submit_async` accept NumPy scalars of any data type whose values are plain bytes, including structured data types and those of `ml_dtypes` such as `bfloat16` and `complex32`, as well as Python `bool`, `int`, `float`, and `complex` scalars, which are passed as `numpy.bool`, `numpy.int64`, `numpy.float64`, and `numpy.complex128` respectively [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) +* Added `DPCTL_BYTES16_T` enumerator to `DPCTLKernelArgType` for 16-byte scalar and local accessor kernel arguments [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) +* Added a `sycl::ext::oneapi::bfloat16` type caster to `dpctl4pybind11.hpp`, defined when the compiler provides the `sycl_ext_oneapi_bfloat16` extension [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) ### Changed +* `dpctl.SyclQueue.copy`, `dpctl.SyclQueue.copy_async`, `dpctl.SyclQueue.fill`, `dpctl.SyclQueue.fill_async`, `dpctl.LocalAccessor`, `dpctl.WorkGroupMemory`, and `dpctl.compiler.SpecializationConstant` accept a data type in any form `numpy.dtype` accepts, for any data type whose values are plain bytes, rather than a fixed set of type strings [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) +* `dpctl.compiler.SpecializationConstant` accepts NumPy scalars and Python `bool`, `int`, `float`, and `complex` scalars [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) +* Added `DPCTL_BYTES16_T` to `DPCTLKernelArgType` before `DPCTL_UNSUPPORTED_KERNEL_ARG`, which shifts the value of the latter [gh-2335](https://github.com/IntelPython/dpctl/pull/2335) ### Deprecated diff --git a/docs/doc_sources/api_reference/dpctl/index.rst b/docs/doc_sources/api_reference/dpctl/index.rst index aa84cb12b9..7244a68dfc 100644 --- a/docs/doc_sources/api_reference/dpctl/index.rst +++ b/docs/doc_sources/api_reference/dpctl/index.rst @@ -77,6 +77,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 +: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 + + 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`` + .. rubric:: Exceptions .. autosummary:: diff --git a/docs/doc_sources/api_reference/dpctl_pybind11.rst b/docs/doc_sources/api_reference/dpctl_pybind11.rst index b2367fec58..9582045f2e 100644 --- a/docs/doc_sources/api_reference/dpctl_pybind11.rst +++ b/docs/doc_sources/api_reference/dpctl_pybind11.rst @@ -44,4 +44,26 @@ We can now use it from Python: # device descriptor print(dev.name) +Type caster for ``bfloat16`` +---------------------------- + +When the DPC++ compiler used to compile an extension provides the +``sycl_ext_oneapi_bfloat16`` extension, ``dpctl4pybind11.hpp`` defines the macro +``DPCTL_HAS_BFLOAT16`` and a type caster connecting +``sycl::ext::oneapi::bfloat16`` to Python. A function taking such an argument +accepts any Python object convertible to a :py:class:`float`, such as an +``ml_dtypes.bfloat16`` scalar, and rounds its value to the nearest +``bfloat16``. A returned ``bfloat16`` becomes a Python :py:class:`float`. + +Since the caster is only defined when the extension is available, code that has +to build with compilers that do not provide it should guard its use: + +.. code-block:: c++ + + #ifdef DPCTL_HAS_BFLOAT16 + sycl::ext::oneapi::bfloat16 round_trip(sycl::ext::oneapi::bfloat16 value) { + return value; + } + #endif + .. _pybind11_url: https://pybind11.readthedocs.io/ diff --git a/docs/doc_sources/known_words.txt b/docs/doc_sources/known_words.txt index b68a926685..b1547b7c65 100644 --- a/docs/doc_sources/known_words.txt +++ b/docs/doc_sources/known_words.txt @@ -82,3 +82,5 @@ profiler keepdims broadcasted broadcastable +bfloat +dtypes diff --git a/dpctl/_backend.pxd b/dpctl/_backend.pxd index 3708ea9130..8deb5d58aa 100644 --- a/dpctl/_backend.pxd +++ b/dpctl/_backend.pxd @@ -73,6 +73,7 @@ cdef extern from "syclinterface/dpctl_sycl_enum_types.h": _LOCAL_ACCESSOR "DPCTL_LOCAL_ACCESSOR", _WORK_GROUP_MEMORY "DPCTL_WORK_GROUP_MEMORY" _RAW_KERNEL_ARG "DPCTL_RAW_KERNEL_ARG" + _BYTES16 "DPCTL_BYTES16_T" ctypedef enum _queue_property_type "DPCTLQueuePropertyType": _DEFAULT_PROPERTY "DPCTL_DEFAULT_PROPERTY" diff --git a/dpctl/_dtype_utils.py b/dpctl/_dtype_utils.py new file mode 100644 index 0000000000..9e73b56ce7 --- /dev/null +++ b/dpctl/_dtype_utils.py @@ -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) + 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 diff --git a/dpctl/_sycl_queue.pxd b/dpctl/_sycl_queue.pxd index c72d09faed..ccc1676506 100644 --- a/dpctl/_sycl_queue.pxd +++ b/dpctl/_sycl_queue.pxd @@ -63,7 +63,8 @@ cdef public api class SyclQueue (_SyclQueue) [ self, list args, void **kargs, - _arg_data_type *kargty + _arg_data_type *kargty, + list _converted ) cdef int _populate_range(self, size_t Range[3], list gS, size_t nGS) @@ -103,13 +104,13 @@ cdef public api class SyclQueue (_SyclQueue) [ cdef DPCTLSyclQueueRef get_queue_ref(self) cpdef memcpy(self, dest, src, size_t count) cpdef SyclEvent memcpy_async(self, dest, src, size_t count, list dEvents=*) - cpdef copy(self, dest, src, size_t count, str dtype=*) + cpdef copy(self, dest, src, size_t count, dtype=*) cpdef SyclEvent copy_async( - self, dest, src, size_t count, list dEvents=*, str dtype=* + self, dest, src, size_t count, list dEvents=*, dtype=* ) - cpdef fill(self, dest, value, size_t count, str dtype=*) + cpdef fill(self, dest, value, size_t count, dtype=*) cpdef SyclEvent fill_async( - self, dest, value, size_t count, list dEvents=*, str dtype=* + self, dest, value, size_t count, list dEvents=*, dtype=* ) cpdef memset(self, mem, int val, size_t count=*) cpdef SyclEvent memset_async( diff --git a/dpctl/_sycl_queue.pyx b/dpctl/_sycl_queue.pyx index 3fbf9c7bf6..cc3b1a10af 100644 --- a/dpctl/_sycl_queue.pyx +++ b/dpctl/_sycl_queue.pyx @@ -85,6 +85,9 @@ from .memory._memory cimport _Memory import ctypes import numbers +import numpy as np + +from ._dtype_utils import as_dtype, as_scalar, is_supported_dtype, pack_value from .enum_types import backend_type from cpython cimport pycapsule @@ -97,13 +100,13 @@ from cpython.buffer cimport ( PyObject_CheckBuffer, PyObject_GetBuffer, ) +from cpython.bytes cimport PyBytes_AS_STRING from cpython.ref cimport Py_INCREF, PyObject from libc.stdint cimport uint8_t, uint16_t, uint32_t, uint64_t from libc.stdlib cimport free, malloc import collections.abc import logging -import struct import sys @@ -125,6 +128,16 @@ __all__ = [ _logger = logging.getLogger(__name__) +# kernel argument type used to pass a value of the given size in bytes +cdef dict _arg_type_by_itemsize = { + 1: _arg_data_type._UINT8_T, + 2: _arg_data_type._UINT16_T, + 4: _arg_data_type._UINT32_T, + 8: _arg_data_type._UINT64_T, + 16: _arg_data_type._BYTES16, +} + + cdef class kernel_arg_type_attribute: cdef str parent_name cdef str attr_name @@ -158,18 +171,11 @@ cdef class LocalAccessor: ``sycl::local_accessor``, to be used as a kernel argument type. Args: - dtype (str): - the data type of the local memory. - The permitted values are - - `'i1'`, `'i2'`, `'i4'`, `'i8'`: - signed integral types int8_t, int16_t, int32_t, int64_t - `'u1'`, `'u2'`, `'u4'`, `'u8'` - unsigned integral types uint8_t, uint16_t, uint32_t, - uint64_t - `'f4'`, `'f8'`, - single- and double-precision floating-point types float and - double + dtype: + the data type of the local memory, anything accepted by + :class:`numpy.dtype`, e.g. `'f4'`, `numpy.float16`, or a + structured data type. Its item size must be 1, 2, 4, 8, or 16 + bytes. shape (tuple, list): Size of LocalAccessor dimensions. Dimension of the LocalAccessor is determined by the length of the tuple. Must be of length 1, 2, or 3, @@ -186,11 +192,11 @@ cdef class LocalAccessor: ValueError: If the shape contains a negative integer. ValueError: - If the dtype string is unrecognized. + If the dtype is unrecognized or of an unsupported size. """ cdef _md_local_accessor lacc - def __cinit__(self, str dtype, shape): + def __cinit__(self, dtype, shape): if not isinstance(shape, (list, tuple)): raise TypeError( f"`shape` must be a list or tuple, got {type(shape)}" @@ -214,28 +220,14 @@ cdef class LocalAccessor: self.lacc.dim1 = shape[1] if ndim > 1 else 1 self.lacc.dim2 = shape[2] if ndim > 2 else 1 - if dtype == "i1": - self.lacc.dpctl_type_id = _arg_data_type._INT8_T - elif dtype == "u1": - self.lacc.dpctl_type_id = _arg_data_type._UINT8_T - elif dtype == "i2": - self.lacc.dpctl_type_id = _arg_data_type._INT16_T - elif dtype == "u2": - self.lacc.dpctl_type_id = _arg_data_type._UINT16_T - elif dtype == "i4": - self.lacc.dpctl_type_id = _arg_data_type._INT32_T - elif dtype == "u4": - self.lacc.dpctl_type_id = _arg_data_type._UINT32_T - elif dtype == "i8": - self.lacc.dpctl_type_id = _arg_data_type._INT64_T - elif dtype == "u8": - self.lacc.dpctl_type_id = _arg_data_type._UINT64_T - elif dtype == "f4": - self.lacc.dpctl_type_id = _arg_data_type._FLOAT - elif dtype == "f8": - self.lacc.dpctl_type_id = _arg_data_type._DOUBLE - else: - raise ValueError(f"Unrecognized type value: '{dtype}'") + dt = as_dtype(dtype) + argty = _arg_type_by_itemsize.get(dt.itemsize) + if argty is None: + raise ValueError( + f"Data type '{dt}' of {dt.itemsize} bytes can not be the type " + "of a local accessor" + ) + self.lacc.dpctl_type_id = <_arg_data_type>argty def __repr__(self): return f"LocalAccessor({self.lacc.ndim})" @@ -394,6 +386,15 @@ cdef class _kernel_arg_type: _arg_data_type._RAW_KERNEL_ARG ) + @property + def dpctl_bytes16(self): + cdef str p_name = "dpctl_bytes16" + return kernel_arg_type_attribute( + self._name, + p_name, + _arg_data_type._BYTES16 + ) + kernel_arg_type = _kernel_arg_type() @@ -488,24 +489,12 @@ ctypedef DPCTLSyclEventRef (*queue_copy_with_events_fn)( ) -cdef size_t _get_dtype_size(str dtype) except *: +cdef size_t _get_dtype_size(object dtype) except *: """ - Parse numpy-style dtype string and return element size in bytes. - Supports: i1, u1, i2, u2, i4, u4, i8, u8, f4, f8 + Return the element size in bytes of a data type accepted by + ``numpy.dtype``. """ - if dtype == "i1" or dtype == "u1": - return 1 - elif dtype == "i2" or dtype == "u2": - return 2 - elif dtype == "i4" or dtype == "u4" or dtype == "f4": - return 4 - elif dtype == "i8" or dtype == "u8" or dtype == "f8": - return 8 - else: - raise ValueError( - f"Unrecognized dtype '{dtype}'. " - "Expected one of: i1, u1, i2, u2, i4, u4, i8, u8, f4, f8" - ) + return (as_dtype(dtype).itemsize) cdef DPCTLSyclEventRef _copy_memcpy_impl( @@ -583,48 +572,19 @@ cdef DPCTLSyclEventRef _copy_memcpy_impl( return ERef -# struct format per dtype used to pack the fill value into its byte pattern. -# The pattern's length is the element size. -# Complex types are packed as (real, imag). -_fill_dtype_formats = { - "i1": "b", - "u1": "B", - "i2": "h", - "u2": "H", - "i4": "i", - "u4": "I", - "f4": "f", - "i8": "q", - "u8": "Q", - "f8": "d", - "c8": "ff", - "c16": "dd", -} - - -cdef bytes _pack_fill_pattern(object value, str dtype): +cdef bytes _pack_fill_pattern(object value, object dtype): """ Return the native-byte-order byte pattern of ``value`` interpreted as a single element of ``dtype``. ``sycl::queue::fill`` replicates this pattern across the destination allocation. """ - cdef str fmt - if dtype not in _fill_dtype_formats: + cdef object dt = as_dtype(dtype) + if dt.itemsize not in (1, 2, 4, 8, 16): raise ValueError( - f"Unrecognized dtype '{dtype}'. Expected one of: " - "i1, u1, i2, u2, i4, u4, i8, u8, f4, f8, c8, c16" - ) - fmt = _fill_dtype_formats[dtype] - try: - if dtype == "c8" or dtype == "c16": - # complex: pack real then imag - 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 '{dtype}': {e}" + f"Fill with dtype '{dt}' of {dt.itemsize} bytes is not supported, " + "the element size must be 1, 2, 4, 8, or 16 bytes" ) + return pack_value(value, dt) cdef DPCTLSyclEventRef _fill_impl( @@ -632,7 +592,7 @@ cdef DPCTLSyclEventRef _fill_impl( object dst, object value, size_t count, - str dtype, + object dtype, DPCTLSyclEventRef *dep_events, size_t dep_events_count ) except *: @@ -734,7 +694,7 @@ cdef DPCTLSyclEventRef _copy_impl( size_t count, DPCTLSyclEventRef *dep_events, size_t dep_events_count, - str dtype="u1" + object dtype="u1" ) except *: # ``count`` is in elements of ``dtype`` (default "u1" => bytes). cdef size_t element_size = _get_dtype_size(dtype) @@ -1181,11 +1141,35 @@ cdef class SyclQueue(_SyclQueue): self, list args, void **kargs, - _arg_data_type *kargty + _arg_data_type *kargty, + list _converted ): cdef int ret = 0 + cdef bytes raw + cdef object argty + cdef object raw_arg for idx, arg in enumerate(args): - if isinstance(arg, ctypes.c_char): + if isinstance(arg, np.generic): + # NumPy scalars are passed as their bytes, which are kept + # alive in ``_converted`` until the arguments are set + dt = arg.dtype + if not is_supported_dtype(dt): + ret = -1 + continue + raw = (arg.tobytes()) + argty = _arg_type_by_itemsize.get(len(raw)) + if argty is not None: + _converted.append(raw) + kargs[idx] = PyBytes_AS_STRING(raw) + kargty[idx] = <_arg_data_type>argty + elif DPCTLRawKernelArg_Available(): + raw_arg = RawKernelArg(raw) + _converted.append(raw_arg) + kargs[idx] = (raw_arg._ref) + kargty[idx] = _arg_data_type._RAW_KERNEL_ARG + else: + ret = -1 + elif isinstance(arg, ctypes.c_char): kargs[idx] = (ctypes.addressof(arg)) kargty[idx] = _arg_data_type._INT8_T elif isinstance(arg, ctypes.c_uint8): @@ -1216,7 +1200,7 @@ cdef class SyclQueue(_SyclQueue): kargs[idx] = (ctypes.addressof(arg)) kargty[idx] = _arg_data_type._DOUBLE elif isinstance(arg, _Memory): - kargs[idx]= (arg._pointer) + kargs[idx] = (arg._pointer) kargty[idx] = _arg_data_type._VOID_PTR elif isinstance(arg, WorkGroupMemory): kargs[idx] = (arg._ref) @@ -1428,7 +1412,11 @@ cdef class SyclQueue(_SyclQueue): kernel (dpctl.compiler.SyclKernel): SYCL kernel object args (List[object]): - List of kernel arguments + List of kernel arguments. A NumPy scalar is passed as its + bytes. A Python ``bool``, ``int``, ``float`` or ``complex`` is + passed as a ``numpy.bool``, ``numpy.int64``, ``numpy.float64`` + or ``numpy.complex128`` respectively, which must match the type + of the kernel parameter. gS (List[int]): Global iteration range. Must be a list of length 1, 2, or 3. lS (List[int], optional): @@ -1472,6 +1460,7 @@ cdef class SyclQueue(_SyclQueue): cdef size_t nLS = len(lS) if lS is not None else 0 cdef size_t nDE = len(dEvents) if dEvents is not None else 0 cdef size_t nKA = 0 + cdef list _converted = [] if kernel.get_kernel_ref() is NULL: raise SyclKernelSubmitError( @@ -1487,6 +1476,9 @@ cdef class SyclQueue(_SyclQueue): f"argument(s), got {len(args)}." ) + # Python scalars are passed as the NumPy scalars they convert to + args = [as_scalar(arg) for arg in args] + # Allocate the arrays to be sent to DPCTLQueue_Submit kargs = malloc(len(args) * sizeof(void*)) if not kargs: @@ -1519,7 +1511,7 @@ cdef class SyclQueue(_SyclQueue): ) # populate the args and argstype arrays - ret = self._populate_args(args, kargs, kargty) + ret = self._populate_args(args, kargs, kargty, _converted) if ret == -1: free(kargs) free(kargty) @@ -1612,7 +1604,11 @@ cdef class SyclQueue(_SyclQueue): kernel (dpctl.compiler.SyclKernel): SYCL kernel object args (List[object]): - List of kernel arguments + List of kernel arguments. A NumPy scalar is passed as its + bytes. A Python ``bool``, ``int``, ``float`` or ``complex`` is + passed as a ``numpy.bool``, ``numpy.int64``, ``numpy.float64`` + or ``numpy.complex128`` respectively, which must match the type + of the kernel parameter. gS (List[int]): Global iteration range. Must be a list of length 1, 2, or 3. lS (List[int], optional): @@ -1696,7 +1692,7 @@ cdef class SyclQueue(_SyclQueue): return SyclEvent._create(ERef) - cpdef copy(self, dest, src, size_t count, str dtype="u1"): + cpdef copy(self, dest, src, size_t count, dtype="u1"): """Copy ``count`` elements of type ``dtype`` from ``src`` to ``dest`` and wait. @@ -1717,11 +1713,11 @@ cdef class SyclQueue(_SyclQueue): protocol. count (int): Number of elements to copy. - dtype (str, optional): - Data type string of the elements to copy. Determines the - element size used to convert ``count`` into a byte count. + dtype (optional): + Data type of the elements to copy, anything accepted by + :class:`numpy.dtype`. Determines the element size used to + convert ``count`` into a byte count. Defaults to ``"u1"`` (one byte per element). - Supported types: i1, u1, i2, u2, i4, u4, i8, u8, f4, f8. """ cdef DPCTLSyclEventRef ERef = NULL @@ -1735,7 +1731,7 @@ cdef class SyclQueue(_SyclQueue): DPCTLEvent_Delete(ERef) cpdef SyclEvent copy_async( - self, dest, src, size_t count, list dEvents=None, str dtype="u1" + self, dest, src, size_t count, list dEvents=None, dtype="u1" ): """Copy ``count`` elements of type ``dtype`` from ``src`` to ``dest`` asynchronously. @@ -1756,11 +1752,11 @@ cdef class SyclQueue(_SyclQueue): Number of elements to copy. dEvents (List[dpctl.SyclEvent], optional): Events that this copy depends on. - dtype (str, optional): - Data type string of the elements to copy. Determines the - element size used to convert ``count`` into a byte count. + dtype (optional): + Data type of the elements to copy, anything accepted by + :class:`numpy.dtype`. Determines the element size used to + convert ``count`` into a byte count. Defaults to ``"u1"`` (one byte per element). - Supported types: i1, u1, i2, u2, i4, u4, i8, u8, f4, f8. Returns: dpctl.SyclEvent: @@ -1798,7 +1794,7 @@ cdef class SyclQueue(_SyclQueue): return SyclEvent._create(ERef) - cpdef fill(self, dest, value, size_t count, str dtype="u1"): + cpdef fill(self, dest, value, size_t count, dtype="u1"): """Fill ``dest`` with ``count`` copies of ``value`` and wait. Internally, this dispatches ``sycl::queue::fill``, which sets each of @@ -1813,19 +1809,18 @@ cdef class SyclQueue(_SyclQueue): Args: dest (dpctl.memory._Memory): Destination USM allocation to fill. - value (int, float, complex): - Value used to fill ``dest``. It is reinterpreted according to - ``dtype``: integral ``dtype`` values expect a Python integer, - ``"f4"`` and ``"f8"`` expect a Python float, and ``"c8"`` and - ``"c16"`` expect a Python complex. + value: + Value used to fill ``dest``, converted to ``dtype``. Integral + ``dtype`` values expect an integer, floating-point ones a real + number and complex ones a complex number. Other data types are + converted by NumPy. count (int): Number of elements to fill. - dtype (str, optional): - Data type string of the fill elements. Determines the element - size and how ``value`` is interpreted. + dtype (optional): + Data type of the fill elements, anything accepted by + :class:`numpy.dtype` of 1, 2, 4, 8, or 16 bytes. Determines + the element size and how ``value`` is interpreted. Defaults to ``"u1"`` (one byte per element). - Supported types: i1, u1, i2, u2, i4, u4, i8, u8, f4, f8, c8, - c16. Raises: TypeError: @@ -1850,7 +1845,7 @@ cdef class SyclQueue(_SyclQueue): DPCTLEvent_Delete(ERef) cpdef SyclEvent fill_async( - self, dest, value, size_t count, list dEvents=None, str dtype="u1" + self, dest, value, size_t count, list dEvents=None, dtype="u1" ): """Fill ``dest`` with ``count`` copies of ``value`` asynchronously. @@ -1868,21 +1863,20 @@ cdef class SyclQueue(_SyclQueue): Args: dest (dpctl.memory._Memory): Destination USM allocation to fill. - value (int, float, complex): - Value used to fill ``dest``. It is reinterpreted according to - ``dtype``: integral ``dtype`` values expect a Python integer, - ``"f4"`` and ``"f8"`` expect a Python float, and ``"c8"`` and - ``"c16"`` expect a Python complex. + value: + Value used to fill ``dest``, converted to ``dtype``. Integral + ``dtype`` values expect an integer, floating-point ones a real + number and complex ones a complex number. Other data types are + converted by NumPy. count (int): Number of elements to fill. dEvents (List[dpctl.SyclEvent], optional): Events that this fill depends on. - dtype (str, optional): - Data type string of the fill elements. Determines the element - size and how ``value`` is interpreted. + dtype (optional): + Data type of the fill elements, anything accepted by + :class:`numpy.dtype` of 1, 2, 4, 8, or 16 bytes. Determines + the element size and how ``value`` is interpreted. Defaults to ``"u1"`` (one byte per element). - Supported types: i1, u1, i2, u2, i4, u4, i8, u8, f4, f8, c8, - c16. Returns: dpctl.SyclEvent: @@ -2316,8 +2310,8 @@ cdef class WorkGroupMemory: memory. - If the constructor is invoked with two arguments, the first argument is - interpreted as the datatype of the local memory, using the numpy type - naming scheme. + interpreted as the datatype of the local memory, anything accepted by + :class:`numpy.dtype`, e.g. ``"f4"``. The second argument is interpreted as the number of elements to allocate. The number of bytes to allocate is then computed from the byte size of the data type and the element count. @@ -2349,26 +2343,18 @@ cdef class WorkGroupMemory: ) nbytes = (args[0]) else: - if not isinstance(args[0], str): - raise TypeError( - "WorkGroupMemory constructor expects first argument to be " - f"`str`, but got {type(args[0])}" - ) if not isinstance(args[1], numbers.Integral): raise TypeError( "WorkGroupMemory constructor expects second argument to " f"be `int`, but got {type(args[1])}" ) - dtype = (args[0]) + dtype = args[0] count = (args[1]) - if not dtype[0] in ["i", "u", "f"]: - raise TypeError(f"Unrecognized type value: '{dtype}'") try: - bit_width = int(dtype[1:]) + byte_size = as_dtype(dtype).itemsize except ValueError: raise TypeError(f"Unrecognized type value: '{dtype}'") - byte_size = bit_width nbytes = count * byte_size self._mem_ref = DPCTLWorkGroupMemory_Create(nbytes) diff --git a/dpctl/apis/include/dpctl4pybind11.hpp b/dpctl/apis/include/dpctl4pybind11.hpp index ecfc85596f..fa8142257e 100644 --- a/dpctl/apis/include/dpctl4pybind11.hpp +++ b/dpctl/apis/include/dpctl4pybind11.hpp @@ -39,6 +39,12 @@ #include #include +#if defined(SYCL_EXT_ONEAPI_BFLOAT16) && __has_include() +#include +// defined when the type caster for sycl::ext::oneapi::bfloat16 is available +#define DPCTL_HAS_BFLOAT16 1 +#endif + namespace py = pybind11; namespace dpctl @@ -607,6 +613,58 @@ template <> struct type_caster PYBIND11_TYPE_CASTER(sycl::half, _("float")); }; +#ifdef DPCTL_HAS_BFLOAT16 +/* This type caster associates the ``sycl::ext::oneapi::bfloat16`` C++ class + * with Python :class:`float` for the purposes of generation of Python bindings + * by pybind11. Any object convertible to :class:`float`, such as an + * ``ml_dtypes.bfloat16`` scalar, is accepted. + */ +template <> struct type_caster +{ +public: + bool load(handle src, bool convert) + { + using bfloat16_t = sycl::ext::oneapi::bfloat16; + + double py_value; + + if (!src) { + return false; + } + + PyObject *source = src.ptr(); + + if (convert || PyFloat_Check(source)) { + py_value = PyFloat_AsDouble(source); + } + else { + return false; + } + + bool py_err = (py_value == double(-1)) && PyErr_Occurred(); + + if (py_err) { + PyErr_Clear(); + if (convert && (PyNumber_Check(source) != 0)) { + auto tmp = reinterpret_steal(PyNumber_Float(source)); + return load(tmp, false); + } + return false; + } + value = bfloat16_t(static_cast(py_value)); + return true; + } + + static handle + cast(sycl::ext::oneapi::bfloat16 src, return_value_policy, handle) + { + return PyFloat_FromDouble(static_cast(static_cast(src))); + } + + PYBIND11_TYPE_CASTER(sycl::ext::oneapi::bfloat16, _("float")); +}; +#endif + } // namespace detail } // namespace pybind11 diff --git a/dpctl/compiler/_compiler.pyx b/dpctl/compiler/_compiler.pyx index 58d7a56310..6a02711c65 100644 --- a/dpctl/compiler/_compiler.pyx +++ b/dpctl/compiler/_compiler.pyx @@ -88,6 +88,8 @@ import numbers import numpy as np +from dpctl._dtype_utils import as_dtype, as_scalar, is_supported_dtype + __all__ = [ "create_kernel_bundle_from_source", "create_kernel_bundle_from_spirv", @@ -360,24 +362,33 @@ cdef class SpecializationConstant: - ``SpecializationConstant(spec_id, obj)`` If the constructor is invoked with a single variadic argument, the - argument is expected to either expose the Python buffer protocol or be - coercible to a NumPy array. If the argument is coercible to a NumPy array - or is one, it must have a supported data type (bool, integral, or - floating point). The specialization constant will be constructed from the - data in the buffer + argument is expected to be a NumPy or Python scalar, or to either expose + the Python buffer protocol or be coercible to a NumPy array. A Python + ``bool``, ``int``, ``float`` or ``complex`` is converted to a + ``numpy.bool``, ``numpy.int64``, ``numpy.float64`` or + ``numpy.complex128`` respectively. If the argument is coercible to a + NumPy array or is one, its data type must be one whose values are plain + bytes, e.g. bool, integral, floating point, structured, or + ``ml_dtypes.bfloat16``. The specialization constant will be constructed + from the data in the buffer - ``SpecializationConstant(spec_id, dtype, obj)`` If the constructor is invoked with two variadic arguments, and the first - argument is a string, it is interpreted as a NumPy ``dtype`` string and - the second argument will be coerced to a NumPy array with that data type. - The data type specified by the first argument must be a supported data - type (bool, integral, or floating point). + argument is not an integer, it is interpreted as a data type, anything + accepted by :class:`numpy.dtype`, and the second argument will be coerced + to a NumPy array with that data type. The same restrictions on the data + type apply. - ``SpecializationConstant(spec_id, nbytes, raw_ptr)`` If the constructor is invoked with two variadic arguments where both are integers, the first argument is interpreted as the number of bytes and the second argument is interpreted as a pointer to the data. + Note that DPC++ splits a specialization constant of a composite type into + one specialization constant per scalar member, each with its own + specialization ID. ``sycl::ext::oneapi::bfloat16`` is such a type, with a + single 16-bit member. + Note that construction of the :class:`.SpecializationConstant` copies the input, so modifications made after construction of the :class:`.SpecializationConstant` will not be reflected in the @@ -401,6 +412,7 @@ cdef class SpecializationConstant: def __cinit__(self, spec_id, *args): cdef int ret_code = 0 cdef object target_obj = None + cdef object target_dtype = None cdef Py_buffer _local_buffer cdef void *copied_data @@ -436,30 +448,34 @@ cdef class SpecializationConstant: target_obj = PyBytes_FromStringAndSize( args[1], args[0] ) - elif isinstance(args[0], str): - target_obj = np.ascontiguousarray(args[1], dtype=args[0]) else: - raise TypeError( - "Invalid arguments." - ) + try: + target_dtype = as_dtype(args[0]) + except ValueError as e: + raise TypeError(f"Invalid arguments: {e}") from None + target_obj = np.ascontiguousarray(args[1], dtype=target_dtype) elif len(args) == 1: - target_obj = args[0] - if not PyObject_CheckBuffer(target_obj): + target_obj = as_scalar(args[0]) + if isinstance(target_obj, np.generic) or ( + not PyObject_CheckBuffer(target_obj) + ): # attempt to coerce to a numpy array target_obj = np.ascontiguousarray(target_obj) if isinstance(target_obj, np.ndarray): - if target_obj.dtype.kind not in ("b", "i", "u", "f", "c"): + target_dtype = target_obj.dtype + if not is_supported_dtype(target_dtype): raise TypeError( "Coercion of input to buffer resulted in an unsupported " - f"data type '{target_obj.dtype}'. When coercing objects, " + f"data type '{target_dtype}'. When coercing objects, " "`SpecializationConstant` expects the data to coerce to a " - "supported type: bool, integral, or real or complex " - "floating point. To pass arbitrary data, use a " - "`memoryview` or `bytes` object, or pass the pointer and " - "size directly." + "data type whose values are plain bytes. To pass " + "arbitrary data, use a `memoryview` or `bytes` object, or " + "pass the pointer and size directly." ) + # not every NumPy data type supports the buffer protocol + target_obj = target_obj.tobytes() ret_code = PyObject_GetBuffer( target_obj, &(_local_buffer), PyBUF_SIMPLE | PyBUF_ANY_CONTIGUOUS diff --git a/dpctl/tests/input_files/bfloat16_spec_const_kernel.cpp b/dpctl/tests/input_files/bfloat16_spec_const_kernel.cpp new file mode 100644 index 0000000000..f4098929c4 --- /dev/null +++ b/dpctl/tests/input_files/bfloat16_spec_const_kernel.cpp @@ -0,0 +1,79 @@ +//===-- bfloat16_spec_const_kernel.cpp - source of a test SPIR-V module ---===// +// +// 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. +// +//===----------------------------------------------------------------------===// +/// +/// \file +/// Source of bfloat16_spec_const_kernel.spv, which is used by +/// dpctl/tests/test_sycl_compiler.py to test that a specialization constant of +/// the type of the DPC++ ``bfloat16`` extension can be set from dpctl. +/// +/// The SPIR-V module was produced with +/// +/// icpx -fsycl -O2 -save-offload-code=. bfloat16_spec_const_kernel.cpp +/// +/// which writes the device image of the kernel next to the executable. +/// +/// The value of the specialization constant is widened to a float by hand, +/// rather than by converting the bfloat16, so that the module does not import +/// the DPC++ bfloat16 devicelib: the SPIR-V interoperability path does not link +/// device libraries, and a module with an unresolved import can not be run. +//===----------------------------------------------------------------------===// + +#include +#include +#include + +using bfloat16_t = sycl::ext::oneapi::bfloat16; + +// the default value has to be the value-initialized bfloat16, as bfloat16 has +// no constexpr constructor +constexpr sycl::specialization_id bf16_spec_const{}; + +class Bfloat16SpecConstKernel; + +int main() +{ + sycl::queue q; + const std::size_t n = 4; + float *x = sycl::malloc_shared(n, q); + float *y = sycl::malloc_shared(n, q); + for (std::size_t i = 0; i < n; ++i) { + x[i] = 1.0f; + y[i] = 0.0f; + } + + q.submit([&](sycl::handler &cgh) { + cgh.set_specialization_constant(bfloat16_t(1.5f)); + cgh.parallel_for( + sycl::range<1>(n), [=](sycl::id<1> i, sycl::kernel_handler kh) { + const float value = x[i]; + const bfloat16_t factor = + kh.get_specialization_constant(); + const std::uint32_t bits = + static_cast( + sycl::bit_cast(factor)) + << 16; + y[i] = value * sycl::bit_cast(bits); + }); + }).wait(); + + sycl::free(x, q); + sycl::free(y, q); + return 0; +} diff --git a/dpctl/tests/input_files/bfloat16_spec_const_kernel.spv b/dpctl/tests/input_files/bfloat16_spec_const_kernel.spv new file mode 100644 index 0000000000..7ee6d9ce4f Binary files /dev/null and b/dpctl/tests/input_files/bfloat16_spec_const_kernel.spv differ diff --git a/dpctl/tests/test_kernel_arg_dtypes.py b/dpctl/tests/test_kernel_arg_dtypes.py new file mode 100644 index 0000000000..ccb1257b4e --- /dev/null +++ b/dpctl/tests/test_kernel_arg_dtypes.py @@ -0,0 +1,425 @@ +# 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. + +"""Tests for NumPy data types as kernel arguments, local accessor and work +group memory types.""" + +import numpy as np +import pytest + +import dpctl +import dpctl.compiler as dpc +import dpctl.memory as dpm + +# layout of std::complex +complex_half = np.dtype([("real", np.float16), ("imag", np.float16)]) + + +def _opencl_queue_or_skip(): + try: + return dpctl.SyclQueue("opencl") + except dpctl.SyclQueueCreationError: + pytest.skip("OpenCL queue could not be created") + + +def _sycl_source_queue_or_skip(): + if not dpc.is_sycl_source_compilation_available(): + pytest.skip("SYCL source compilation extension not available") + for filter_str in ("level_zero", "opencl:gpu", "opencl:cpu"): + try: + q = dpctl.SyclQueue(filter_str) + except dpctl.SyclQueueCreationError: + continue + if q.sycl_device.can_compile("sycl"): + return q + pytest.skip("No device supporting SYCL source compilation is available") + + +def _skip_if_no_aspect(q, ctype): + if "half" in ctype and not q.sycl_device.has_aspect_fp16: + pytest.skip("Device does not support half precision floating point") + if "double" in ctype and not q.sycl_device.has_aspect_fp64: + pytest.skip("Device does not support double precision floating point") + + +def _pragmas(ctype): + src = "" + if "half" in ctype: + src += "#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n" + if "double" in ctype: + src += "#pragma OPENCL EXTENSION cl_khr_fp64 : enable\n" + return src + + +def _to_device(q, host): + usm = dpm.MemoryUSMDevice(host.nbytes, queue=q) + q.memcpy(dest=usm, src=host, count=host.nbytes) + return usm + + +@pytest.mark.parametrize( + "ctype,dtype", + [ + ("char", "i1"), + ("uchar", "u1"), + ("short", "i2"), + ("ushort", "u2"), + ("int", "i4"), + ("uint", "u4"), + ("long", "i8"), + ("ulong", "u8"), + ("half", "f2"), + ("float", "f4"), + ("double", "f8"), + ], +) +def test_numpy_scalar_arg(ctype, dtype): + q = _opencl_queue_or_skip() + _skip_if_no_aspect(q, ctype) + src = _pragmas(ctype) + ( + f"kernel void axpy(global {ctype} *a, global {ctype} *b, {ctype} d) {{" + " size_t i = get_global_id(0);" + " b[i] = d * a[i] + b[i];" + "}" + ) + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("axpy") + + n = 64 + a = (np.arange(n) % 5).astype(dtype) + b = np.ones(n, dtype=dtype) + a_usm, b_usm = _to_device(q, a), _to_device(q, b) + + d = np.dtype(dtype).type(3) + q.submit(krn, [a_usm, b_usm, d], [n]).wait() + q.memcpy(dest=b, src=b_usm, count=b.nbytes) + assert np.array_equal(b, (a * d + np.ones(n, dtype=dtype)).astype(dtype)) + + +def _complex_half_scalar(kind, value): + if kind == "complex32": + ml_dtypes = pytest.importorskip("ml_dtypes") + return ml_dtypes.complex32(value) + return np.array((value.real, value.imag), dtype=complex_half)[()] + + +@pytest.mark.parametrize( + "ctype,kind", + [ + ("half2", "complex32"), + ("half2", "structured"), + ("float2", "c8"), + ("double2", "c16"), + ], +) +def test_complex_scalar_arg(ctype, kind): + q = _opencl_queue_or_skip() + _skip_if_no_aspect(q, ctype) + src = _pragmas(ctype) + ( + f"kernel void cmul(global {ctype} *a, global {ctype} *b, {ctype} d) {{" + " size_t i = get_global_id(0);" + f" {ctype} ai = a[i];" + f" b[i] = ({ctype})(ai.x * d.x - ai.y * d.y, ai.x * d.y + ai.y * d.x);" + "}" + ) + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("cmul") + + n = 64 + a = (np.arange(n) % 4) - 1j * (np.arange(n) % 3) + d = 2 - 1j + expected = a * d + + if ctype == "half2": + # host arrays are kept as float16 pairs + a_host = a.astype("c8").view("f4").astype("f2") + d_arg = _complex_half_scalar(kind, d) + else: + a_host = a.astype(kind) + d_arg = np.dtype(kind).type(d) + b_host = np.zeros_like(a_host) + a_usm, b_usm = _to_device(q, a_host), _to_device(q, b_host) + + if ctype == "half2": + assert d_arg.dtype.itemsize == 2 * a_host.itemsize + q.submit(krn, [a_usm, b_usm, d_arg], [n]).wait() + q.memcpy(dest=b_host, src=b_usm, count=b_host.nbytes) + + if ctype == "half2": + res = b_host.astype("f4").view("c8") + else: + res = b_host + assert np.array_equal(res, expected.astype(res.dtype)) + + +def test_bcomplex32_scalar_arg(): + ml_dtypes = pytest.importorskip("ml_dtypes") + q = _opencl_queue_or_skip() + src = ( + "kernel void widen(global float *out, ushort2 d) {" + " out[0] = as_float(((uint)d.x) << 16);" + " out[1] = as_float(((uint)d.y) << 16);" + "}" + ) + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("widen") + out = np.zeros(2, dtype="f4") + out_usm = _to_device(q, out) + q.submit(krn, [out_usm, ml_dtypes.bcomplex32(1.5 - 2.25j)], [1]).wait() + q.memcpy(dest=out, src=out_usm, count=out.nbytes) + assert out.tolist() == [1.5, -2.25] + + +@pytest.mark.parametrize( + "ctype,value", + [("uchar", True), ("long", -(2**40)), ("double", 0.1), ("double2", 1 - 2j)], +) +def test_python_scalar_arg(ctype, value): + q = _opencl_queue_or_skip() + _skip_if_no_aspect(q, ctype) + src = _pragmas(ctype) + ( + f"kernel void put(global {ctype} *out, {ctype} d) {{ out[0] = d; }}" + ) + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("put") + expected = { + bool: np.bool_, + int: np.int64, + float: np.float64, + complex: np.complex128, + }[type(value)](value) + out = np.zeros(1, dtype=expected.dtype) + out_usm = _to_device(q, out) + q.submit(krn, [out_usm, value], [1]).wait() + q.memcpy(dest=out, src=out_usm, count=out.nbytes) + assert out[0] == expected + + +def test_python_int_arg_out_of_range(): + q = _opencl_queue_or_skip() + src = "kernel void put(global long *out, long d) { out[0] = d; }" + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("put") + out_usm = dpm.MemoryUSMDevice(8, queue=q) + with pytest.raises(OverflowError): + q.submit(krn, [out_usm, 2**63], [1]) + + +def test_bfloat16_scalar_arg(): + ml_dtypes = pytest.importorskip("ml_dtypes") + q = _opencl_queue_or_skip() + # a bfloat16 is the upper half of a float with the same value + src = ( + "kernel void scale(global float *a, global float *b, ushort d) {" + " size_t i = get_global_id(0);" + " b[i] = a[i] * as_float(((uint)d) << 16);" + "}" + ) + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("scale") + + n = 64 + a = np.arange(n, dtype="f4") + b = np.zeros_like(a) + a_usm, b_usm = _to_device(q, a), _to_device(q, b) + + d = ml_dtypes.bfloat16(1.5) + q.submit(krn, [a_usm, b_usm, d], [n]).wait() + q.memcpy(dest=b, src=b_usm, count=b.nbytes) + assert np.array_equal(b, a * np.float32(d)) + + +def test_bfloat16_sycl_source_arg(): + ml_dtypes = pytest.importorskip("ml_dtypes") + q = _sycl_source_queue_or_skip() + src = """ + #include + #include + + namespace syclext = sycl::ext::oneapi::experimental; + using bfloat16_t = sycl::ext::oneapi::bfloat16; + + extern "C" SYCL_EXTERNAL + SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclext::nd_range_kernel<1>)) + void scale_bfloat16(bfloat16_t *a, bfloat16_t *b, bfloat16_t d) { + sycl::nd_item<1> it = + sycl::ext::oneapi::this_work_item::get_nd_item<1>(); + size_t i = it.get_global_linear_id(); + float p = static_cast(a[i]) * static_cast(d); + // truncated on the bits, since float to bfloat16 conversion gives + // wrong results on some Intel GPUs; the products are exact + b[i] = sycl::bit_cast( + std::uint16_t(sycl::bit_cast(p) >> 16)); + } + """ + try: + kb = dpc.create_kernel_bundle_from_sycl_source(q, src) + except Exception: + pytest.skip(f"Compilation of source failed for {q.sycl_device}") + krn = kb.get_sycl_kernel("scale_bfloat16") + + lws = 32 + n = lws * 4 + a = (0.5 * (np.arange(n) % 7)).astype(ml_dtypes.bfloat16) + # bfloat16 arrays do not support the buffer protocol + a_bits = a.view(np.uint16) + b_bits = np.zeros_like(a_bits) + a_usm, b_usm = _to_device(q, a_bits), _to_device(q, b_bits) + + d = ml_dtypes.bfloat16(3.0) + q.submit(krn, [a_usm, b_usm, d], [n], [lws]).wait() + q.memcpy(dest=b_bits, src=b_usm, count=b_bits.nbytes) + assert np.array_equal(b_bits.view(ml_dtypes.bfloat16), a * d) + + +def test_structured_scalar_raw_kernel_arg(): + if not dpctl.RawKernelArg.is_available(): + pytest.skip("Raw kernel arg extension not available") + q = _sycl_source_queue_or_skip() + src = """ + #include + + namespace syclext = sycl::ext::oneapi::experimental; + + struct Coeffs { + float scale; + float offset; + int repeats; + }; + + extern "C" SYCL_EXTERNAL + SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclext::nd_range_kernel<1>)) + void apply_coeffs(float *a, float *b, Coeffs c) { + sycl::nd_item<1> it = + sycl::ext::oneapi::this_work_item::get_nd_item<1>(); + size_t i = it.get_global_linear_id(); + float acc = a[i]; + for (int r = 0; r < c.repeats; ++r) { + acc = acc * c.scale + c.offset; + } + b[i] = acc; + } + """ + try: + kb = dpc.create_kernel_bundle_from_sycl_source(q, src) + except Exception: + pytest.skip(f"Compilation of source failed for {q.sycl_device}") + krn = kb.get_sycl_kernel("apply_coeffs") + + coeffs_t = np.dtype([("scale", "f4"), ("offset", "f4"), ("repeats", "i4")]) + assert coeffs_t.itemsize == 12 + + lws = 16 + n = lws * 4 + a = np.arange(n, dtype="f4") + b = np.zeros_like(a) + a_usm, b_usm = _to_device(q, a), _to_device(q, b) + + c = np.array((2.0, 1.0, 3), dtype=coeffs_t)[()] + q.submit(krn, [a_usm, b_usm, c], [n], [lws]).wait() + q.memcpy(dest=b, src=b_usm, count=b.nbytes) + + expected = a.copy() + for _ in range(3): + expected = expected * np.float32(2.0) + np.float32(1.0) + assert np.array_equal(b, expected) + + +@pytest.mark.parametrize( + "ctype,dtype,base", + [ + ("half2", complex_half, "f2"), + ("float2", np.dtype("c8"), "f4"), + ("long2", np.dtype([("x", "i8"), ("y", "i8")]), "i8"), + ], +) +def test_local_accessor_arg(ctype, dtype, base): + q = _opencl_queue_or_skip() + _skip_if_no_aspect(q, ctype) + src = _pragmas(ctype) + ( + f"kernel void rev(global {ctype} *a, local {ctype} *tmp) {{" + " size_t lid = get_local_id(0);" + " size_t lws = get_local_size(0);" + " tmp[lid] = a[get_global_id(0)];" + " barrier(CLK_LOCAL_MEM_FENCE);" + " a[get_global_id(0)] = tmp[lws - 1 - lid];" + "}" + ) + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("rev") + + lws = 16 + n = lws * 4 + # host arrays are kept in the base type of the data type + x = np.arange(2 * n, dtype=base) + x_usm = _to_device(q, x) + + la = dpctl.LocalAccessor(dtype, (lws,)) + q.submit(krn, [x_usm, la], [n], [lws]).wait() + res = np.empty_like(x) + q.memcpy(dest=res, src=x_usm, count=res.nbytes) + + expected = x.reshape(n // lws, lws, 2)[:, ::-1, :].reshape(-1) + assert np.array_equal(res, expected) + + +@pytest.mark.parametrize( + "dtype", + ["i1", "u2", "f2", "f4", "c8", "c16", "float", np.float16, complex_half], +) +def test_local_accessor_dtypes(dtype): + assert isinstance(dpctl.LocalAccessor(dtype, (8,)), dpctl.LocalAccessor) + + +def test_local_accessor_bfloat16(): + ml_dtypes = pytest.importorskip("ml_dtypes") + for dtype in (ml_dtypes.bfloat16, "bfloat16", ml_dtypes.complex32): + la = dpctl.LocalAccessor(dtype, (8,)) + assert isinstance(la, dpctl.LocalAccessor) + + +@pytest.mark.parametrize( + "dtype", + [ + "not_a_type", + "U4", + "S4", + "M8", + "f16", + object, + [("a", "f4"), ("b", "f4"), ("c", "f4")], + ], +) +def test_local_accessor_invalid_dtypes(dtype): + with pytest.raises(ValueError): + dpctl.LocalAccessor(dtype, (8,)) + + +def test_work_group_memory_dtypes(): + if not dpctl.WorkGroupMemory.is_available(): + pytest.skip("Workgroup memory extension not available") + for dtype in ("f4", np.float16, complex_half, "c16"): + assert dpctl.WorkGroupMemory(dtype, 8)._ref != 0 + # any size is allowed + assert dpctl.WorkGroupMemory([("a", "f4"), ("b", "i8")], 8)._ref != 0 + for dtype in ("not_a_type", "U4", object): + with pytest.raises(TypeError): + dpctl.WorkGroupMemory(dtype, 8) + + +@pytest.mark.parametrize( + "arg", [np.str_("ab"), np.datetime64("2026-01-01"), "ab", None] +) +def test_unsupported_scalar_args(arg): + q = _opencl_queue_or_skip() + src = "kernel void k(global int *a, int d) { a[get_global_id(0)] = d; }" + krn = dpc.create_kernel_bundle_from_source(q, src).get_sycl_kernel("k") + a_usm = dpm.MemoryUSMDevice(4 * 8, queue=q) + with pytest.raises(TypeError): + q.submit(krn, [a_usm, arg], [8]) diff --git a/dpctl/tests/test_sycl_compiler.py b/dpctl/tests/test_sycl_compiler.py index 45775a3268..421e69d19c 100644 --- a/dpctl/tests/test_sycl_compiler.py +++ b/dpctl/tests/test_sycl_compiler.py @@ -369,6 +369,109 @@ def test_create_kernel_bundle_with_composite_spec_const(): assert np.all(y == 12.5) +def test_create_kernel_bundle_with_bfloat16_spec_const(): + ml_dtypes = pytest.importorskip("ml_dtypes") + try: + q = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Could not create default queue") + + factor = ml_dtypes.bfloat16(1.5) + sp = dpc.SpecializationConstant(0, factor) + + spirv_file = get_spirv_abspath("bfloat16_spec_const_kernel.spv") + with open(spirv_file, "br") as spv: + spv_bytes = spv.read() + + kb = dpc.create_kernel_bundle_from_spirv(q, spv_bytes, specializations=[sp]) + kernel = kb.get_sycl_kernel("_ZTS23Bfloat16SpecConstKernel") + + n = 128 + x = np.arange(1, n + 1, dtype="f4") + y = np.zeros_like(x) + + x_usm = dpctl.memory.MemoryUSMDevice(x.nbytes, queue=q) + y_usm = dpctl.memory.MemoryUSMDevice(y.nbytes, queue=q) + + e1 = q.memcpy_async(x_usm, x, x.nbytes) + e2 = q.submit(kernel, [x_usm, y_usm], [n], dEvents=[e1]) + e3 = q.memcpy_async(y, y_usm, y.nbytes, [e2]) + + ht_e = q._submit_keep_args_alive([x_usm], [e3]) + + e3.wait() + ht_e.wait() + + assert np.array_equal(y, x * np.float32(factor)) + + +def test_spec_const_from_numpy_dtypes(): + ml_dtypes = pytest.importorskip("ml_dtypes") + bf16 = dpc.SpecializationConstant(0, ml_dtypes.bfloat16(1.5)) + assert bf16.size == 2 + assert dpc.SpecializationConstant(0, "bfloat16", 1.5) == bf16 + assert dpc.SpecializationConstant(0, ml_dtypes.bfloat16, 1.5) == bf16 + assert dpc.SpecializationConstant(0, np.uint16(0x3FC0)) == bf16 + + assert dpc.SpecializationConstant(0, np.int32(42)) == ( + dpc.SpecializationConstant(0, "i4", 42) + ) + assert dpc.SpecializationConstant(0, np.float32, 2.5) == ( + dpc.SpecializationConstant(0, "f4", 2.5) + ) + + +def test_spec_const_complex_half(): + complex_half = np.dtype([("real", "f2"), ("imag", "f2")]) + sp = dpc.SpecializationConstant(0, complex_half, (1.0, -2.0)) + assert sp.size == 4 + assert sp == dpc.SpecializationConstant( + 0, np.array([1.0, -2.0], dtype="f2") + ) + assert sp == dpc.SpecializationConstant( + 0, np.array((1.0, -2.0), dtype=complex_half)[()] + ) + + +def test_spec_const_complex32(): + ml_dtypes = pytest.importorskip("ml_dtypes") + sp = dpc.SpecializationConstant(0, ml_dtypes.complex32(1 - 2j)) + assert sp.size == 4 + assert sp == dpc.SpecializationConstant(0, "complex32", 1 - 2j) + assert sp == dpc.SpecializationConstant( + 0, np.array([1.0, -2.0], dtype="f2") + ) + + +@pytest.mark.parametrize( + "value,dtype", + [(True, "?"), (42, "i8"), (2.5, "f8"), (1 - 2j, "c16")], +) +def test_spec_const_from_python_scalar(value, dtype): + sp = dpc.SpecializationConstant(0, value) + assert sp == dpc.SpecializationConstant(0, dtype, value) + + +def test_spec_const_invalid_dtypes(): + with pytest.raises(TypeError): + dpc.SpecializationConstant(0, object(), 42) + with pytest.raises(TypeError): + dpc.SpecializationConstant(0, "not_a_data_type", 42) + with pytest.raises(TypeError): + dpc.SpecializationConstant(0, "U4", "abc") + + +def test_spirv_specializations_parser_bfloat16(): + spirv_file = get_spirv_abspath("bfloat16_spec_const_kernel.spv") + with open(spirv_file, "rb") as spv: + spv_bytes = spv.read() + spec_consts = parse_spirv_specializations(spv_bytes) + assert len(spec_consts) == 1 + # the parser reports the 16-bit member of the bfloat16 + assert spec_consts[0].dtype == "u2" + assert spec_consts[0].itemsize == 2 + + def test_specialization_constant_addressof_ref(): sp = dpc.SpecializationConstant(0, "i4", 42) ref = sp.addressof_ref() diff --git a/dpctl/tests/test_sycl_kernel_submit.py b/dpctl/tests/test_sycl_kernel_submit.py index e173035ae6..ef204f398b 100644 --- a/dpctl/tests/test_sycl_kernel_submit.py +++ b/dpctl/tests/test_sycl_kernel_submit.py @@ -303,6 +303,7 @@ def test_kernel_arg_type(): _check_kernel_arg_type_instance(kernel_arg_type.dpctl_local_accessor) _check_kernel_arg_type_instance(kernel_arg_type.dpctl_work_group_memory) _check_kernel_arg_type_instance(kernel_arg_type.dpctl_raw_kernel_arg) + _check_kernel_arg_type_instance(kernel_arg_type.dpctl_bytes16) def get_spirv_abspath(fn): diff --git a/dpctl/tests/test_sycl_queue_copy.py b/dpctl/tests/test_sycl_queue_copy.py index c584ba26b9..496f0f1efd 100644 --- a/dpctl/tests/test_sycl_queue_copy.py +++ b/dpctl/tests/test_sycl_queue_copy.py @@ -223,7 +223,7 @@ def test_copy_with_invalid_dtype(): src = dpctl.memory.MemoryUSMShared(nbytes, queue=q) dst = dpctl.memory.MemoryUSMShared(nbytes, queue=q) - invalid_dtypes = ["i3", "f16", "u16", "x4", "42", "float", ""] + invalid_dtypes = ["i3", "f16", "u16", "x4", "42", "", "U4", "O"] for bad_dtype in invalid_dtypes: with pytest.raises(ValueError) as cm: @@ -296,3 +296,49 @@ def test_copy_without_dtype_backward_compat(): dst_mv = memoryview(dst) assert dst_mv[:nbytes].tobytes() == src_mv[:nbytes].tobytes() + + +@pytest.mark.parametrize( + "dtype", + [ + "float32", + np.float32, + np.dtype("f4"), + "f2", + "c8", + [("re", "f2"), ("im", "f2")], + [("a", "f4"), ("b", "i4"), ("c", "i4")], + ], +) +def test_copy_with_numpy_dtypes(dtype): + """Test typed copy with data types in any form numpy.dtype accepts.""" + try: + q = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Default constructor for SyclQueue failed") + + num_elements = 6 + itemsize = np.dtype(dtype).itemsize + src = np.arange(num_elements * itemsize, dtype=np.uint8) + dst = np.zeros_like(src) + + # only the first half of the elements is copied + q.copy(dst, src, num_elements // 2, dtype=dtype) + half = num_elements // 2 * itemsize + assert np.array_equal(dst[:half], src[:half]) + assert not dst[half:].any() + + +def test_copy_with_ml_dtypes(): + ml_dtypes = pytest.importorskip("ml_dtypes") + try: + q = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Default constructor for SyclQueue failed") + + src = np.arange(8, dtype=np.uint16) + for dtype in (ml_dtypes.bfloat16, "bfloat16", ml_dtypes.complex32): + dst = np.zeros_like(src) + n = src.nbytes // np.dtype(dtype).itemsize + q.copy(dst, src, n, dtype=dtype) + assert np.array_equal(dst, src) diff --git a/dpctl/tests/test_sycl_queue_fill.py b/dpctl/tests/test_sycl_queue_fill.py index 76546858c0..e80662f56f 100644 --- a/dpctl/tests/test_sycl_queue_fill.py +++ b/dpctl/tests/test_sycl_queue_fill.py @@ -263,7 +263,7 @@ def test_fill_with_invalid_dtype(): mem = dpctl.memory.MemoryUSMShared(64, queue=q) - for bad_dtype in ["i3", "f16", "u16", "x4", "42", "float", ""]: + for bad_dtype in ["i3", "f16", "u16", "x4", "42", "", "U4", "O", "V12"]: with pytest.raises(ValueError) as cm: q.fill(mem, 0, 8, dtype=bad_dtype) assert ( @@ -330,3 +330,53 @@ def test_fill_count_out_of_bounds(): q.fill(mem, 0, 9, dtype="i8") with pytest.raises(ValueError): q.fill_async(mem, 0, nbytes + 1) + + +@pytest.mark.parametrize( + "dtype,value", + [ + ("float32", 1.5), + (np.float32, 1.5), + ("f2", -2.5), + ("b1", True), + ("c8", 1 - 2j), + ([("re", "f2"), ("im", "f2")], (1.0, -2.0)), + ([("a", "i8"), ("b", "i8")], (3, -4)), + ], +) +def test_fill_with_numpy_dtypes(dtype, value): + try: + q = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Default constructor for SyclQueue failed") + + num_elements = 8 + np_dt = np.dtype(dtype) + nbytes = num_elements * np_dt.itemsize + mem = dpctl.memory.MemoryUSMShared(nbytes, queue=q) + q.fill(mem, value, num_elements, dtype=dtype) + + expected = np.array([value] * num_elements, dtype=np_dt).tobytes() + assert memoryview(mem)[:nbytes].tobytes() == expected + + +def test_fill_with_ml_dtypes(): + ml_dtypes = pytest.importorskip("ml_dtypes") + try: + q = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Default constructor for SyclQueue failed") + + num_elements = 8 + for dtype, value in ( + (ml_dtypes.bfloat16, 1.5), + ("bfloat16", 1.5), + (ml_dtypes.complex32, 1 - 2j), + (ml_dtypes.bcomplex32, 1 - 2j), + ): + np_dt = np.dtype(dtype) + nbytes = num_elements * np_dt.itemsize + mem = dpctl.memory.MemoryUSMShared(nbytes, queue=q) + q.fill(mem, value, num_elements, dtype=dtype) + expected = np.array([value] * num_elements, dtype=np_dt).tobytes() + assert memoryview(mem)[:nbytes].tobytes() == expected diff --git a/examples/pybind11/use_dpctl_sycl_queue/tests/test_queue_device.py b/examples/pybind11/use_dpctl_sycl_queue/tests/test_queue_device.py index e495c87cd4..831ec1fccf 100644 --- a/examples/pybind11/use_dpctl_sycl_queue/tests/test_queue_device.py +++ b/examples/pybind11/use_dpctl_sycl_queue/tests/test_queue_device.py @@ -17,10 +17,17 @@ # coding: utf-8 import numpy as np +import pytest import use_queue_device as uqd import dpctl +skip_no_bfloat16 = pytest.mark.skipif( + not uqd.has_bfloat16, + reason="The DPC++ compiler the extension was built with does not provide " + "the bfloat16 extension", +) + def test_compute_units(): q = dpctl.SyclQueue() @@ -71,3 +78,33 @@ def test_get_partition_max_sub_devices(): assert type(mt) is int assert mt >= 0 assert mt <= d.max_compute_units + + +@skip_no_bfloat16 +def test_bfloat16_round_trip(): + # a Python float is rounded to the nearest bfloat16, and comes back as the + # float that bfloat16 stands for + assert uqd.bfloat16_round_trip(1.5) == 1.5 + assert uqd.bfloat16_round_trip(3.14) == 3.140625 + assert type(uqd.bfloat16_round_trip(3.14)) is float + # an integer converts, too + assert uqd.bfloat16_round_trip(3) == 3.0 + + +@skip_no_bfloat16 +def test_bfloat16_round_trip_ml_dtypes_scalar(): + ml_dtypes = pytest.importorskip("ml_dtypes") + value = ml_dtypes.bfloat16(3.14) + assert uqd.bfloat16_round_trip(value) == float(value) + + +@skip_no_bfloat16 +def test_offloaded_bfloat16_scale(): + ml_dtypes = pytest.importorskip("ml_dtypes") + execution_queue = dpctl.SyclQueue() + X = np.arange(1, 11, dtype=np.float32) + factor = ml_dtypes.bfloat16(1.1) + + Y = uqd.offloaded_bfloat16_scale(execution_queue, X, factor) + + assert np.array_equal(Y, X * np.float32(factor)) diff --git a/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/__init__.py b/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/__init__.py index 461fc2ebda..3c985fbea2 100644 --- a/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/__init__.py +++ b/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/__init__.py @@ -17,12 +17,15 @@ # coding: utf-8 from ._use_queue_device import ( + bfloat16_round_trip, get_device_global_mem_size, get_device_local_mem_size, get_max_compute_units, get_partition_max_sub_devices, get_sub_group_sizes, + has_bfloat16, offloaded_array_mod, + offloaded_bfloat16_scale, ) __all__ = [ @@ -32,6 +35,9 @@ "offloaded_array_mod", "get_sub_group_sizes", "get_partition_max_sub_devices", + "has_bfloat16", + "bfloat16_round_trip", + "offloaded_bfloat16_scale", ] __doc__ = """ @@ -39,7 +45,8 @@ SYCL entities. dpctl provides type casters that bind ``sycl::queue`` to `dpctl.SyclQueue`, -``sycl::device`` to `dpctl.SyclDevice`, etc. +``sycl::device`` to `dpctl.SyclDevice`, ``sycl::ext::oneapi::bfloat16`` to +Python `float`, etc. Use of these type casters simplifies writing of Python extensions and compile then using SYCL C++ compilers, such as Intel(R) oneAPI DPC++ compiler. diff --git a/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/_example.cpp b/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/_example.cpp index a82aa591b4..68da3aa1ce 100644 --- a/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/_example.cpp +++ b/examples/pybind11/use_dpctl_sycl_queue/use_queue_device/_example.cpp @@ -85,6 +85,64 @@ offloaded_array_mod(sycl::queue q, return res; } +#ifdef DPCTL_HAS_BFLOAT16 +using bfloat16_t = sycl::ext::oneapi::bfloat16; + +bfloat16_t bfloat16_round_trip(bfloat16_t value) { return value; } + +py::array_t +offloaded_bfloat16_scale(sycl::queue q, + py::array_t array, + bfloat16_t factor) +{ + py::buffer_info arg_pybuf = array.request(); + if (arg_pybuf.ndim != 1) { + throw std::runtime_error("Expecting a vector"); + } + + size_t n = arg_pybuf.size; + + auto res = py::array_t(n); + py::buffer_info res_pybuf = res.request(); + + float *a = static_cast(arg_pybuf.ptr); + float *r = static_cast(res_pybuf.ptr); + + { + const sycl::property_list props = { + sycl::property::buffer::use_host_ptr()}; + sycl::buffer a_buf(a, sycl::range<1>(n), props); + sycl::buffer r_buf(r, sycl::range<1>(n), props); + + q.submit([&](sycl::handler &cgh) { + sycl::accessor a_acc(a_buf, cgh, sycl::read_only); + sycl::accessor r_acc(r_buf, cgh, sycl::write_only, sycl::no_init); + + cgh.parallel_for(sycl::range<1>(n), [=](sycl::id<1> idx) { + // the product is not rounded to a bfloat16 on the device, + // since float to bfloat16 conversion gives wrong results on + // some Intel GPUs + r_acc[idx] = a_acc[idx] * static_cast(factor); + }); + }).wait_and_throw(); + } + + return res; +} +#else +py::object bfloat16_round_trip(py::object) +{ + throw std::runtime_error("The DPC++ compiler this extension was built with " + "does not provide the bfloat16 extension"); +} + +py::object offloaded_bfloat16_scale(sycl::queue, py::object, py::object) +{ + throw std::runtime_error("The DPC++ compiler this extension was built with " + "does not provide the bfloat16 extension"); +} +#endif + std::vector get_sub_group_sizes(const sycl::device &d) { return d.get_info(); @@ -116,4 +174,15 @@ PYBIND11_MODULE(_use_queue_device, m, py::mod_gil_not_used()) m.def("get_partition_max_sub_devices", &get_partition_max_sub_devices, "Gets info::device::partition_max_sub_devices for given " "dpctl.SyclDevice"); +#ifdef DPCTL_HAS_BFLOAT16 + m.attr("has_bfloat16") = true; +#else + m.attr("has_bfloat16") = false; +#endif + m.def("bfloat16_round_trip", &bfloat16_round_trip, + "Converts a Python float, or an ml_dtypes.bfloat16 scalar, to a " + "sycl::ext::oneapi::bfloat16 and back"); + m.def("offloaded_bfloat16_scale", &offloaded_bfloat16_scale, + "Compute offloaded scaling of a NumPy array of float32 by a " + "bfloat16 factor"); } diff --git a/libsyclinterface/include/syclinterface/dpctl_sycl_enum_types.h b/libsyclinterface/include/syclinterface/dpctl_sycl_enum_types.h index c66b8330f9..de0580c64a 100644 --- a/libsyclinterface/include/syclinterface/dpctl_sycl_enum_types.h +++ b/libsyclinterface/include/syclinterface/dpctl_sycl_enum_types.h @@ -102,6 +102,7 @@ typedef enum DPCTL_LOCAL_ACCESSOR, DPCTL_WORK_GROUP_MEMORY, DPCTL_RAW_KERNEL_ARG, + DPCTL_BYTES16_T, /* any trivially copyable 16-byte value */ DPCTL_UNSUPPORTED_KERNEL_ARG } DPCTLKernelArgType; diff --git a/libsyclinterface/source/dpctl_sycl_queue_interface.cpp b/libsyclinterface/source/dpctl_sycl_queue_interface.cpp index a7d2a49c8b..e441789ee6 100644 --- a/libsyclinterface/source/dpctl_sycl_queue_interface.cpp +++ b/libsyclinterface/source/dpctl_sycl_queue_interface.cpp @@ -36,6 +36,7 @@ #include #include +#include #include #include #include @@ -113,6 +114,12 @@ using namespace sycl; CGH.set_arg(IDX, la); \ return true; \ } \ + case DPCTL_BYTES16_T: \ + { \ + auto la = local_accessor(R, CGH); \ + CGH.set_arg(IDX, la); \ + return true; \ + } \ default: \ error_handler("Kernel argument could not be created.", __FILE__, \ __func__, __LINE__, error_level::error); \ @@ -216,6 +223,13 @@ bool set_kernel_arg(handler &cgh, case DPCTL_FLOAT64_T: cgh.set_arg(idx, *(double *)Arg); break; + case DPCTL_BYTES16_T: + { + complexNumber Val; + std::memcpy(&Val, Arg, sizeof(Val)); + cgh.set_arg(idx, Val); + break; + } case DPCTL_VOID_PTR: cgh.set_arg(idx, Arg); break;