8367996402

Committed 21 Mar 2024 12:55AM UTC coverage: 82.668% (-0.4%) from 83.041%

Build # 8367996402

Build Type

Pull #1398

github

Committed by

web-flow

Commit Message

Merge 180755115 into 7ef52bf2c

Pull Request Pull Request #1398: Feature/specialization device func

Run Details

1633 of 2249 branches covered (72.61%)

Branch coverage included in aggregate %.

11 of 11 new or added lines in 1 file covered. (100.0%)

29 existing lines in 2 files now uncovered.

6628 of 7744 relevant lines covered (85.59%)

0.86 hits per line

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

74.16

/numba_dpex/decorators.py

# SPDX-FileCopyrightText: 2020 - 2024 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 import config
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 get_compiler
from numba_dpex.core.targets.dpjit_target import DPEX_TARGET_NAME


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 and kws["nopython"] is not True:
        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["pipeline_class"]

    use_mlir = kws.pop("use_mlir", bool(config.USE_MLIR))

    kws.update({"nopython": True})
    kws.update({"parallel": True})
    kws.update({"pipeline_class": get_compiler(use_mlir)})

    kws.update({"_target": DPEX_TARGET_NAME})

    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_TARGET_NAME]] = dpjit

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

IntelPython / numba-dpex / 8367996402

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