8369878729

Committed 21 Mar 2024 04:55AM UTC coverage: 79.854% (-0.02%) from 79.874%

Build # 8369878729

Build Type

Pull #1400

github

Committed by

web-flow

Commit Message

Merge 10d30811c into 98377191f

Pull Request Pull Request #1400: Feature/move experimental decorators to stable

Run Details

1530 of 2238 branches covered (68.36%)

Branch coverage included in aggregate %.

4 of 4 new or added lines in 3 files covered. (100.0%)

37 existing lines in 5 files now uncovered.

6441 of 7744 relevant lines covered (83.17%)

0.83 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.91

/numba_dpex/experimental/decorators.py

# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
#
# SPDX-License-Identifier: Apache-2.0

""" The set of experimental decorators provided by numba_dpex that are not yet
ready to move to numba_dpex.core.
"""
import inspect
from warnings import warn

from numba.core import sigutils, typeinfer
from numba.core.target_extension import (
    jit_registry,
    resolve_dispatcher_from_str,
    target_registry,
)

from numba_dpex.kernel_api_impl.spirv.dispatcher import SPIRVKernelDispatcher
from numba_dpex.kernel_api_impl.spirv.target import CompilationMode

from .target import DPEX_KERNEL_EXP_TARGET_NAME


def _parse_func_or_sig(signature_or_function):
    # Handle signature (borrowed from numba). swapped signature and list check
    if signature_or_function is None:
        # No signature, no function
        pyfunc = None
        sigs = []
    elif sigutils.is_signature(signature_or_function):
        # A single signature is passed
        pyfunc = None
        sigs = [signature_or_function]
    elif isinstance(signature_or_function, list):
        # A list of signatures is passed
        pyfunc = None
        sigs = signature_or_function
    else:
        # A function is passed
        pyfunc = signature_or_function
        sigs = []

    return pyfunc, sigs


def kernel(function_or_signature=None, **options):
    """A decorator to compile a function written using :py:mod:`numba_dpex.kernel_api`.

    The ``kernel`` decorator triggers the compilation of a function written
    using the data-parallel kernel programming API exposed by
    :py:mod:`numba_dpex.kernel_api`. Such a function is conceptually
    equivalent to a kernel function written in the C++ SYCL eDSL. The
    decorator will compile the function based on the types of the arguments
    to a SPIR-V binary that can be executed either on OpenCL CPU, GPU
    devices or Intel Level Zero GPU devices.

    Any function to be compilable using the kernel decorator should
    adhere to the following semantic rules:

    - The first argument to the function should be either an instance of the
      :class:`numba_dpex.kernel_api.Item` class or an instance of the
      :class:`numba_dpex.kernel_api.NdItem`.

    - The function should not return any value.

    - The function should have at least one array type argument that can
      either be an instance of ``dpnp.ndarray`` or an instance of
      ``dpctl.tensor.usm_ndarray``.


    Args:
        signature_or_function (optional): An optional signature or list of
            signatures for which a function is to be compiled. Passing in a
            signature "specializes" the decorated function and no other versions of
            the function will be compiled. A function can also be
            directly passed instead of a signature and the signature will get
            inferred from the function. The actual compilation happens on every
            invocation of the :func:`numba_dpex.experimental.call_kernel` function
            where the decorated function is passed in as an argument along with the
            argument values for the decorated function.
        options (optional):
            - **debug** (bool): Whether the compilation should happen in debug
              mode. *(Default = False)*
            - **inline_threshold** (int): Specifies the level of inlining that
              the compiler should attempt. *(Default = 2)*
    Returns:
        An instance of
        :class:`numba_dpex.kernel_api_impl.spirv.dispatcher.KernelDispatcher`.
        The ``KernelDispatcher`` object compiles the decorated function when
        passed in to :func:`numba_dpex.experimental.call_kernel`.

    Examples:

    1. Decorate a function and pass it to ``call_kernel`` for compilation and
       execution.

    .. code-block:: python

        import dpnp
        import numba_dpex as dpex
        from numba_dpex import kernel_api as kapi


        # Data parallel kernel implementing vector sum
        @dpex.kernel
        def vecadd(item: kapi.Item, a, b, c):
            i = item.get_id(0)
            c[i] = a[i] + b[i]


        N = 1024
        a = dpnp.ones(N)
        b = dpnp.ones_like(a)
        c = dpnp.zeros_like(a)
        dpex.call_kernel(vecadd, kapi.Range(N), a, b, c)

    2. Specializes a kernel and then compiles it directly before executing it
       via ``call_kernel``. The kernel is specialized to expect a 1-D
       ``dpnp.ndarray`` with either ``float32`` type elements or ``int64`` type
       elements.

    .. code-block:: python

        import dpnp
        import numba_dpex as dpex
        from numba_dpex import kernel_api as kapi
        from numba_dpex import DpnpNdArray, float32, int64
        from numba_dpex.core.types.kernel_api.index_space_ids import ItemType

        i64arrty = DpnpNdArray(ndim=1, dtype=int64, layout="C")
        f32arrty = DpnpNdArray(ndim=1, dtype=float32, layout="C")
        item_ty = ItemType(ndim=1)

        specialized_kernel = dpex.kernel(
            [
                (item_ty, i64arrty, i64arrty, i64arrty),
                (item_ty, f32arrty, f32arrty, f32arrty),
            ]
        )


        def vecadd(item: kapi.Item, a, b, c):
            i = item.get_id(0)
            c[i] = a[i] + b[i]


        # Compile all specializations for vecadd
        precompiled_kernels = specialized_kernel(vecadd)
        N = 1024
        a = dpnp.ones(N, dtype=dpnp.int64)
        b = dpnp.ones_like(a)
        c = dpnp.zeros_like(a)
        # Call a specific pre-compiled version of vecadd
        dpex.call_kernel(precompiled_kernels, kapi.Range(N), a, b, c)

    """

    # dispatcher is a type:
    # <class 'numba_dpex.experimental.kernel_dispatcher.KernelDispatcher'>
    dispatcher = resolve_dispatcher_from_str(DPEX_KERNEL_EXP_TARGET_NAME)
    if "_compilation_mode" in options:
        user_compilation_mode = options["_compilation_mode"]
        warn(
            "_compilation_mode is an internal flag that should not be set "
            "in the decorator. The decorator defined option "
            f"{user_compilation_mode} is going to be ignored."
        )
    options["_compilation_mode"] = CompilationMode.KERNEL

    # TODO: The options need to be evaluated and checked here like it is
    # done in numba.core.decorators.jit

    func, sigs = _parse_func_or_sig(function_or_signature)
    for sig in sigs:
        if isinstance(sig, str):
            raise NotImplementedError(
                "Specifying signatures as string is not yet supported"
            )

    def _kernel_dispatcher(pyfunc):
        disp: SPIRVKernelDispatcher = dispatcher(
            pyfunc=pyfunc,
            targetoptions=options,
        )

        if len(sigs) > 0:
            with typeinfer.register_dispatcher(disp):
                for sig in sigs:
                    disp.compile(sig)
                disp.disable_compile()

        return disp

    if func is None:
        return _kernel_dispatcher

    if not inspect.isfunction(func):
        raise ValueError(
            "Argument passed to the kernel decorator is neither a "
            "function object, nor a signature. If you are trying to "
            "specialize the kernel that takes a single argument, specify "
            "the return type as None explicitly."
        )
    return _kernel_dispatcher(func)


def device_func(func_or_sig=None, **options):
    """Compiles a function into a "device-only" function.

    The decorator is used to compile overloads in the DpexKernelTarget and
    users should use the decorator to define functions that are only callable
    from inside another device_func or a kernel.

    A device_func is not compiled down to device binary IR and instead left as
    LLVM IR. It is done so that the function can be inlined fully into the
    kernel module from where it is used at the LLVM level, leading to more
    optimization opportunities.

    Returns:
        KernelDispatcher: A KernelDispatcher instance with the
        _compilation_mode option set to DEVICE_FUNC.
    """
    dispatcher = resolve_dispatcher_from_str(DPEX_KERNEL_EXP_TARGET_NAME)

    if "_compilation_mode" in options:
        user_compilation_mode = options["_compilation_mode"]
        warn(
            "_compilation_mode is an internal flag that should not be set "
            "in the decorator. The decorator defined option "
            f"{user_compilation_mode} is going to be ignored."
        )
    options["_compilation_mode"] = CompilationMode.DEVICE_FUNC

    func, sigs = _parse_func_or_sig(func_or_sig)
    for sig in sigs:
        if isinstance(sig, str):
            raise NotImplementedError(
                "Specifying signatures as string is not yet supported"
            )

    def _kernel_dispatcher(pyfunc):
        disp: SPIRVKernelDispatcher = dispatcher(
            pyfunc=pyfunc,
            targetoptions=options,
        )

        if len(sigs) > 0:
            with typeinfer.register_dispatcher(disp):
                for sig in sigs:
                    disp.compile(sig)
                disp.disable_compile()

        return disp

    if func is None:
        return _kernel_dispatcher

    return _kernel_dispatcher(func_or_sig)


jit_registry[target_registry[DPEX_KERNEL_EXP_TARGET_NAME]] = device_func

1	# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	""" The set of experimental decorators provided by numba_dpex that are not yet	1✔
6	ready to move to numba_dpex.core.
7	"""
8	import inspect	1✔
9	from warnings import warn	1✔
10
11	from numba.core import sigutils, typeinfer	1✔
12	from numba.core.target_extension import (	1✔
13	jit_registry,
14	resolve_dispatcher_from_str,
15	target_registry,
16	)
17
18	from numba_dpex.kernel_api_impl.spirv.dispatcher import SPIRVKernelDispatcher	1✔
19	from numba_dpex.kernel_api_impl.spirv.target import CompilationMode	1✔
20
21	from .target import DPEX_KERNEL_EXP_TARGET_NAME	1✔
22
23
24	def _parse_func_or_sig(signature_or_function):	1✔
25	# Handle signature (borrowed from numba). swapped signature and list check
26	if signature_or_function is None:	1✔
27	# No signature, no function
28	pyfunc = None	1✔
29	sigs = []	1✔
30	elif sigutils.is_signature(signature_or_function):	1✔
31	# A single signature is passed
32	pyfunc = None	1✔
33	sigs = [signature_or_function]	1✔
34	elif isinstance(signature_or_function, list):	1✔
35	# A list of signatures is passed
36	pyfunc = None	1✔
37	sigs = signature_or_function	1✔
38	else:
39	# A function is passed
40	pyfunc = signature_or_function	1✔
41	sigs = []	1✔
42
43	return pyfunc, sigs	1✔
44
45
46	def kernel(function_or_signature=None, **options):	1✔
47	"""A decorator to compile a function written using :py:mod:`numba_dpex.kernel_api`.
48
49	The ``kernel`` decorator triggers the compilation of a function written
50	using the data-parallel kernel programming API exposed by
51	:py:mod:`numba_dpex.kernel_api`. Such a function is conceptually
52	equivalent to a kernel function written in the C++ SYCL eDSL. The
53	decorator will compile the function based on the types of the arguments
54	to a SPIR-V binary that can be executed either on OpenCL CPU, GPU
55	devices or Intel Level Zero GPU devices.
56
57	Any function to be compilable using the kernel decorator should
58	adhere to the following semantic rules:
59
60	- The first argument to the function should be either an instance of the
61	:class:`numba_dpex.kernel_api.Item` class or an instance of the
62	:class:`numba_dpex.kernel_api.NdItem`.
63
64	- The function should not return any value.
65
66	- The function should have at least one array type argument that can
67	either be an instance of ``dpnp.ndarray`` or an instance of
68	``dpctl.tensor.usm_ndarray``.
69
70
71	Args:
72	signature_or_function (optional): An optional signature or list of
73	signatures for which a function is to be compiled. Passing in a
74	signature "specializes" the decorated function and no other versions of
75	the function will be compiled. A function can also be
76	directly passed instead of a signature and the signature will get
77	inferred from the function. The actual compilation happens on every
78	invocation of the :func:`numba_dpex.experimental.call_kernel` function
79	where the decorated function is passed in as an argument along with the
80	argument values for the decorated function.
81	options (optional):
82	- debug (bool): Whether the compilation should happen in debug
83	mode. (Default = False)
84	- inline_threshold (int): Specifies the level of inlining that
85	the compiler should attempt. (Default = 2)
86	Returns:
87	An instance of
88	:class:`numba_dpex.kernel_api_impl.spirv.dispatcher.KernelDispatcher`.
89	The ``KernelDispatcher`` object compiles the decorated function when
90	passed in to :func:`numba_dpex.experimental.call_kernel`.
91
92	Examples:
93
94	1. Decorate a function and pass it to ``call_kernel`` for compilation and
95	execution.
96
97	.. code-block:: python
98
99	import dpnp
100	import numba_dpex as dpex
101	from numba_dpex import kernel_api as kapi
102
103
104	# Data parallel kernel implementing vector sum
105	@dpex.kernel
106	def vecadd(item: kapi.Item, a, b, c):
107	i = item.get_id(0)
108	c[i] = a[i] + b[i]
109
110
111	N = 1024
112	a = dpnp.ones(N)
113	b = dpnp.ones_like(a)
114	c = dpnp.zeros_like(a)
115	dpex.call_kernel(vecadd, kapi.Range(N), a, b, c)
116
117	2. Specializes a kernel and then compiles it directly before executing it
118	via ``call_kernel``. The kernel is specialized to expect a 1-D
119	``dpnp.ndarray`` with either ``float32`` type elements or ``int64`` type
120	elements.
121
122	.. code-block:: python
123
124	import dpnp
125	import numba_dpex as dpex
126	from numba_dpex import kernel_api as kapi
127	from numba_dpex import DpnpNdArray, float32, int64
128	from numba_dpex.core.types.kernel_api.index_space_ids import ItemType
129
130	i64arrty = DpnpNdArray(ndim=1, dtype=int64, layout="C")
131	f32arrty = DpnpNdArray(ndim=1, dtype=float32, layout="C")
132	item_ty = ItemType(ndim=1)
133
134	specialized_kernel = dpex.kernel(
135	[
136	(item_ty, i64arrty, i64arrty, i64arrty),
137	(item_ty, f32arrty, f32arrty, f32arrty),
138	]
139	)
140
141
142	def vecadd(item: kapi.Item, a, b, c):
143	i = item.get_id(0)
144	c[i] = a[i] + b[i]
145
146
147	# Compile all specializations for vecadd
148	precompiled_kernels = specialized_kernel(vecadd)
149	N = 1024
150	a = dpnp.ones(N, dtype=dpnp.int64)
151	b = dpnp.ones_like(a)
152	c = dpnp.zeros_like(a)
153	# Call a specific pre-compiled version of vecadd
154	dpex.call_kernel(precompiled_kernels, kapi.Range(N), a, b, c)
155
156	"""
157
158	# dispatcher is a type:
159	# <class 'numba_dpex.experimental.kernel_dispatcher.KernelDispatcher'>
160	dispatcher = resolve_dispatcher_from_str(DPEX_KERNEL_EXP_TARGET_NAME)	1✔
161	if "_compilation_mode" in options:	1✔
162	user_compilation_mode = options["_compilation_mode"]	1✔
163	warn(	1✔
164	"_compilation_mode is an internal flag that should not be set "
165	"in the decorator. The decorator defined option "
166	f"{user_compilation_mode} is going to be ignored."
167	)
168	options["_compilation_mode"] = CompilationMode.KERNEL	1✔
169
170	# TODO: The options need to be evaluated and checked here like it is
171	# done in numba.core.decorators.jit
172
173	func, sigs = _parse_func_or_sig(function_or_signature)	1✔
174	for sig in sigs:	1✔
175	if isinstance(sig, str):	1✔
176	raise NotImplementedError(
177	"Specifying signatures as string is not yet supported"
178	)
179
180	def _kernel_dispatcher(pyfunc):	1✔
181	disp: SPIRVKernelDispatcher = dispatcher(	1✔
182	pyfunc=pyfunc,
183	targetoptions=options,
184	)
185
186	if len(sigs) > 0:	1✔
187	with typeinfer.register_dispatcher(disp):	1✔
188	for sig in sigs:	1✔
189	disp.compile(sig)	1✔
190	disp.disable_compile()	1✔
191
192	return disp	1✔
193
194	if func is None:	1✔
195	return _kernel_dispatcher	1✔
196
197	if not inspect.isfunction(func):	1✔
198	raise ValueError(	1✔
199	"Argument passed to the kernel decorator is neither a "
200	"function object, nor a signature. If you are trying to "
201	"specialize the kernel that takes a single argument, specify "
202	"the return type as None explicitly."
203	)
204	return _kernel_dispatcher(func)	1✔
205
206
207	def device_func(func_or_sig=None, **options):	1✔
208	"""Compiles a function into a "device-only" function.
209
210	The decorator is used to compile overloads in the DpexKernelTarget and
211	users should use the decorator to define functions that are only callable
212	from inside another device_func or a kernel.
213
214	A device_func is not compiled down to device binary IR and instead left as
215	LLVM IR. It is done so that the function can be inlined fully into the
216	kernel module from where it is used at the LLVM level, leading to more
217	optimization opportunities.
218
219	Returns:
220	KernelDispatcher: A KernelDispatcher instance with the
221	_compilation_mode option set to DEVICE_FUNC.
222	"""
223	dispatcher = resolve_dispatcher_from_str(DPEX_KERNEL_EXP_TARGET_NAME)	1✔
224
225	if "_compilation_mode" in options:	1!
UNCOV 226	user_compilation_mode = options["_compilation_mode"]	×
UNCOV 227	warn(	×
228	"_compilation_mode is an internal flag that should not be set "
229	"in the decorator. The decorator defined option "
230	f"{user_compilation_mode} is going to be ignored."
231	)
232	options["_compilation_mode"] = CompilationMode.DEVICE_FUNC	1✔
233
234	func, sigs = _parse_func_or_sig(func_or_sig)	1✔
235	for sig in sigs:	1✔
236	if isinstance(sig, str):	1✔
237	raise NotImplementedError(
238	"Specifying signatures as string is not yet supported"
239	)
240
241	def _kernel_dispatcher(pyfunc):	1✔
242	disp: SPIRVKernelDispatcher = dispatcher(	1✔
243	pyfunc=pyfunc,
244	targetoptions=options,
245	)
246
247	if len(sigs) > 0:	1✔
248	with typeinfer.register_dispatcher(disp):	1✔
249	for sig in sigs:	1✔
250	disp.compile(sig)	1✔
251	disp.disable_compile()	1✔
252
253	return disp	1✔
254
255	if func is None:	1✔
256	return _kernel_dispatcher	1✔
257
258	return _kernel_dispatcher(func_or_sig)	1✔
259
260
261	jit_registry[target_registry[DPEX_KERNEL_EXP_TARGET_NAME]] = device_func	1✔

IntelPython / numba-dpex / 8369878729

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous