5662065863

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

import inspect
import warnings

from numba.core import decorators, sigutils
from numba.core.target_extension import jit_registry, target_registry

from numba_dpex.core.kernel_interface.dispatcher import JitKernel
from numba_dpex.core.kernel_interface.func import (
    compile_func,
    compile_func_template,
)
from numba_dpex.core.pipelines.dpjit_compiler import DpjitCompiler


def kernel(
    func_or_sig=None,
    debug=False,
    enable_cache=True,
):
    """A decorator to define a kernel function.

    A kernel function is conceptually equivalent to a SYCL kernel function, and
    gets compiled into either an OpenCL or a LevelZero SPIR-V binary kernel.
    A kernel decorated Python function has the following restrictions:

        * The function can not return any value.
        * All array arguments passed to a kernel should adhere to compute
          follows data programming model.
    """

    def _kernel_dispatcher(pyfunc, sigs=None):
        return JitKernel(
            pyfunc=pyfunc,
            debug_flags=debug,
            enable_cache=enable_cache,
            specialization_sigs=sigs,
        )

    if func_or_sig is None:
        return _kernel_dispatcher
    elif isinstance(func_or_sig, str):
        raise NotImplementedError(
            "Specifying signatures as string is not yet supported by numba-dpex"
        )
    elif isinstance(func_or_sig, list) or sigutils.is_signature(func_or_sig):
        # String signatures are not supported as passing usm_ndarray type as
        # a string is not possible. Numba's sigutils relies on the type being
        # available in Numba's `types.__dict__` and dpex types are not registered
        # there yet.
        if isinstance(func_or_sig, list):
            for sig in func_or_sig:
                if isinstance(sig, str):
                    raise NotImplementedError(
                        "Specifying signatures as string is not yet supported "
                        "by numba-dpex"
                    )
        # Specialized signatures can either be a single signature or a list.
        # In case only one signature is provided convert it to a list
        if not isinstance(func_or_sig, list):
            func_or_sig = [func_or_sig]

        def _specialized_kernel_dispatcher(pyfunc):
            return JitKernel(
                pyfunc=pyfunc,
                debug_flags=debug,
                enable_cache=enable_cache,
                specialization_sigs=func_or_sig,
            )

        return _specialized_kernel_dispatcher
    else:
        func = func_or_sig
        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 void explicitly."
            )
        return _kernel_dispatcher(func)


def func(func_or_sig=None, debug=False, enable_cache=True):
    """A decorator to define a kernel device function.

    Device functions are functions that can be only invoked from a kernel
    and not from a host function. This provides a special decorator
    `numba_dpex.func` specifically to implement a device function.

    A device function can be invoked from another device function and
    unlike a kernel function, a device function can return a value like
    normal functions.
    """

    def _func_autojit(pyfunc):
        return compile_func_template(
            pyfunc, debug=debug, enable_cache=enable_cache
        )

    if func_or_sig is None:
        return _func_autojit
    elif isinstance(func_or_sig, str):
        raise NotImplementedError(
            "Specifying signatures as string is not yet supported by numba-dpex"
        )
    elif isinstance(func_or_sig, list) or sigutils.is_signature(func_or_sig):
        # String signatures are not supported as passing usm_ndarray type as
        # a string is not possible. Numba's sigutils relies on the type being
        # available in Numba's types.__dict__ and dpex types are not registered
        # there yet.
        if isinstance(func_or_sig, list):
            for sig in func_or_sig:
                if isinstance(sig, str):
                    raise NotImplementedError(
                        "Specifying signatures as string is not yet supported "
                        "by numba-dpex"
                    )
        # Specialized signatures can either be a single signature or a list.
        # In case only one signature is provided convert it to a list
        if not isinstance(func_or_sig, list):
            func_or_sig = [func_or_sig]

        def _wrapped(pyfunc):
            return compile_func(pyfunc, func_or_sig, debug=debug)

        return _wrapped
    else:
        # no signature
        func = func_or_sig
        return _func_autojit(func)


# ----------------- Experimental dpjit decorator ------------------------------#


def dpjit(*args, **kws):
    if "nopython" in kws:
        warnings.warn(
            "nopython is set for dpjit and is ignored", RuntimeWarning
        )
    if "forceobj" in kws:
        warnings.warn(
            "forceobj is set for dpjit and is ignored", RuntimeWarning
        )
        del kws["forceobj"]
    if "pipeline_class" in kws:
        warnings.warn(
            "pipeline class is set for dpjit and is ignored", RuntimeWarning
        )
        del kws["forceobj"]
    kws.update({"nopython": True})
    kws.update({"parallel": True})
    kws.update({"pipeline_class": DpjitCompiler})

    # FIXME: When trying to use dpex's target context, overloads do not work
    # properly. We will turn on dpex target once the issue is fixed.

    # kws.update({"_target": "dpex"})

    return decorators.jit(*args, **kws)


# add it to the decorator registry, this is so e.g. @overload can look up a
# JIT function to do the compilation work.
jit_registry[target_registry["dpex"]] = dpjit

1	# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import inspect	1✔
6	import warnings	1✔
7
8	from numba.core import decorators, sigutils	1✔
9	from numba.core.target_extension import jit_registry, target_registry	1✔
10
11	from numba_dpex.core.kernel_interface.dispatcher import JitKernel	1✔
12	from numba_dpex.core.kernel_interface.func import (	1✔
13	compile_func,
14	compile_func_template,
15	)
16	from numba_dpex.core.pipelines.dpjit_compiler import DpjitCompiler	1✔
17
18
19	def kernel(	1✔
20	func_or_sig=None,
21	debug=False,
22	enable_cache=True,
23	):
24	"""A decorator to define a kernel function.
25
26	A kernel function is conceptually equivalent to a SYCL kernel function, and
27	gets compiled into either an OpenCL or a LevelZero SPIR-V binary kernel.
28	A kernel decorated Python function has the following restrictions:
29
30	* The function can not return any value.
31	* All array arguments passed to a kernel should adhere to compute
32	follows data programming model.
33	"""
34
35	def _kernel_dispatcher(pyfunc, sigs=None):	1✔
36	return JitKernel(	1✔
37	pyfunc=pyfunc,
38	debug_flags=debug,
39	enable_cache=enable_cache,
40	specialization_sigs=sigs,
41	)
42
43	if func_or_sig is None:	1✔
44	return _kernel_dispatcher	1✔
45	elif isinstance(func_or_sig, str):	1✔
46	raise NotImplementedError(
47	"Specifying signatures as string is not yet supported by numba-dpex"
48	)
49	elif isinstance(func_or_sig, list) or sigutils.is_signature(func_or_sig):	1✔
50	# String signatures are not supported as passing usm_ndarray type as
51	# a string is not possible. Numba's sigutils relies on the type being
52	# available in Numba's `types.__dict__` and dpex types are not registered
53	# there yet.
54	if isinstance(func_or_sig, list):	1✔
55	for sig in func_or_sig:	1✔
56	if isinstance(sig, str):	1✔
57	raise NotImplementedError(
58	"Specifying signatures as string is not yet supported "
59	"by numba-dpex"
60	)
61	# Specialized signatures can either be a single signature or a list.
62	# In case only one signature is provided convert it to a list
63	if not isinstance(func_or_sig, list):	1✔
64	func_or_sig = [func_or_sig]	1✔
65
66	def _specialized_kernel_dispatcher(pyfunc):	1✔
67	return JitKernel(	1✔
68	pyfunc=pyfunc,
69	debug_flags=debug,
70	enable_cache=enable_cache,
71	specialization_sigs=func_or_sig,
72	)
73
74	return _specialized_kernel_dispatcher	1✔
75	else:
76	func = func_or_sig	1✔
77	if not inspect.isfunction(func):	1✔
78	raise ValueError(	1✔
79	"Argument passed to the kernel decorator is neither a "
80	"function object, nor a signature. If you are trying to "
81	"specialize the kernel that takes a single argument, specify "
82	"the return type as void explicitly."
83	)
84	return _kernel_dispatcher(func)	1✔
85
86
87	def func(func_or_sig=None, debug=False, enable_cache=True):	1✔
88	"""A decorator to define a kernel device function.
89
90	Device functions are functions that can be only invoked from a kernel
91	and not from a host function. This provides a special decorator
92	`numba_dpex.func` specifically to implement a device function.
93
94	A device function can be invoked from another device function and
95	unlike a kernel function, a device function can return a value like
96	normal functions.
97	"""
98
99	def _func_autojit(pyfunc):	1✔
100	return compile_func_template(	1✔
101	pyfunc, debug=debug, enable_cache=enable_cache
102	)
103
104	if func_or_sig is None:	1✔
105	return _func_autojit	1✔
106	elif isinstance(func_or_sig, str):	1✔
107	raise NotImplementedError(
108	"Specifying signatures as string is not yet supported by numba-dpex"
109	)
110	elif isinstance(func_or_sig, list) or sigutils.is_signature(func_or_sig):	1✔
111	# String signatures are not supported as passing usm_ndarray type as
112	# a string is not possible. Numba's sigutils relies on the type being
113	# available in Numba's types.__dict__ and dpex types are not registered
114	# there yet.
115	if isinstance(func_or_sig, list):	1✔
116	for sig in func_or_sig:	1✔
117	if isinstance(sig, str):	1✔
118	raise NotImplementedError(
119	"Specifying signatures as string is not yet supported "
120	"by numba-dpex"
121	)
122	# Specialized signatures can either be a single signature or a list.
123	# In case only one signature is provided convert it to a list
124	if not isinstance(func_or_sig, list):	1✔
125	func_or_sig = [func_or_sig]	1✔
126
127	def _wrapped(pyfunc):	1✔
128	return compile_func(pyfunc, func_or_sig, debug=debug)	1✔
129
130	return _wrapped	1✔
131	else:
132	# no signature
133	func = func_or_sig	1✔
134	return _func_autojit(func)	1✔
135
136
137	# ----------------- Experimental dpjit decorator ------------------------------#
138
139
140	def dpjit(args, *kws):	1✔
141	if "nopython" in kws:	1!
UNCOV 142	warnings.warn(	×
143	"nopython is set for dpjit and is ignored", RuntimeWarning
144	)
145	if "forceobj" in kws:	1!
UNCOV 146	warnings.warn(	×
147	"forceobj is set for dpjit and is ignored", RuntimeWarning
148	)
UNCOV 149	del kws["forceobj"]	×
150	if "pipeline_class" in kws:	1!
UNCOV 151	warnings.warn(	×
152	"pipeline class is set for dpjit and is ignored", RuntimeWarning
153	)
UNCOV 154	del kws["forceobj"]	×
155	kws.update({"nopython": True})	1✔
156	kws.update({"parallel": True})	1✔
157	kws.update({"pipeline_class": DpjitCompiler})	1✔
158
159	# FIXME: When trying to use dpex's target context, overloads do not work
160	# properly. We will turn on dpex target once the issue is fixed.
161
162	# kws.update({"_target": "dpex"})
163
164	return decorators.jit(args, *kws)	1✔
165
166
167	# add it to the decorator registry, this is so e.g. @overload can look up a
168	# JIT function to do the compilation work.
169	jit_registry[target_registry["dpex"]] = dpjit	1✔

IntelPython / numba-dpex / 5662065863

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