8386358402

Committed 22 Mar 2024 06:22AM UTC coverage: 79.059% (-0.8%) from 79.874%

Build # 8386358402

Build Type

Pull #1400

github

Committed by

web-flow

Commit Message

Merge cb01231e0 into 611a59711

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

Run Details

1503 of 2209 branches covered (68.04%)

Branch coverage included in aggregate %.

79 of 86 new or added lines in 3 files covered. (91.86%)

84 existing lines in 12 files now uncovered.

6331 of 7700 relevant lines covered (82.22%)

0.82 hits per line

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

0.0

/numba_dpex/core/kernel_interface/func.py

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

from numba.core import sigutils, types
from numba.core.typing.templates import AbstractTemplate, ConcreteTemplate

from numba_dpex.core import config
from numba_dpex.core.caching import LRUCache, NullCache
from numba_dpex.core.compiler import compile_with_dpex
from numba_dpex.core.descriptor import dpex_kernel_target
from numba_dpex.core.utils import (
    build_key,
    create_func_hash,
    strip_usm_metadata,
)


class DpexFunction(object):
    """Class to materialize dpex function

    Helper class to eager compile a specialized `numba_dpex.func`
    decorated Python function into a LLVM function with `spir_func`
    calling convention.

    A specialized `numba_dpex.func` decorated Python function is one
    where the user has specified a signature or a list of signatures
    for the function. The function gets compiled as soon as the Python
    program is loaded, i.e., eagerly, instead of JIT compilation once
    the function is invoked.
    """

    def __init__(self, pyfunc, debug=False):
        """Constructor for `DpexFunction`

        Args:
            pyfunc (`function`): A python function to be compiled.
            debug (`bool`, optional): Debug option for compilation.
                Defaults to `False`.
        """
        self._pyfunc = pyfunc
        self._debug = debug

    def compile(self, arg_types, return_types):
        """The actual compilation function.

        Args:
            arg_types (`tuple`): Function argument types in a tuple.
            return_types (`numba.core.types.scalars.Integer`):
                An integer value to specify the return type.

        Returns:
            `numba.core.compiler.CompileResult`: The compiled result
        """

        cres = compile_with_dpex(
            pyfunc=self._pyfunc,
            pyfunc_name=self._pyfunc.__name__,
            return_type=return_types,
            target_context=dpex_kernel_target.target_context,
            typing_context=dpex_kernel_target.typing_context,
            args=arg_types,
            is_kernel=False,
            debug=self._debug,
        )
        func = cres.library.get_function(cres.fndesc.llvm_func_name)
        cres.target_context.set_spir_func_calling_conv(func)

        return cres


class DpexFunctionTemplate(object):
    """Helper class to compile an unspecialized `numba_dpex.func`

    A helper class to JIT compile an unspecialized `numba_dpex.func`
    decorated Python function into an LLVM function with `spir_func`
    calling convention.
    """

    def __init__(self, pyfunc, debug=False, enable_cache=True):
        """Constructor for `DpexFunctionTemplate`

        Args:
            pyfunc (function): A python function to be compiled.
            debug (bool, optional): Debug option for compilation.
                Defaults to `False`.
            enable_cache (bool, optional): Flag to turn on/off caching.
                Defaults to `True`.
        """
        self._pyfunc = pyfunc
        self._debug = debug
        self._enable_cache = enable_cache

        self._func_hash = create_func_hash(pyfunc)

        if not config.ENABLE_CACHE:
            self._cache = NullCache()
        elif self._enable_cache:
            self._cache = LRUCache(
                name="DpexFunctionTemplateCache",
                capacity=config.CACHE_SIZE,
                pyfunc=self._pyfunc,
            )
        else:
            self._cache = NullCache()
        self._cache_hits = 0

    @property
    def cache(self):
        """Cache accessor"""
        return self._cache

    @property
    def cache_hits(self):
        """Cache hit count accessor"""
        return self._cache_hits

    def compile(self, args):
        """Compile a `numba_dpex.func` decorated function

        Compile a `numba_dpex.func` decorated Python function with the
        given argument types. Each signature is compiled once by caching
        the compiled function inside this object.

        Args:
            args (`tuple`): Function argument types in a tuple.

        Returns:
            `numba.core.typing.templates.Signature`: Signature of the
                compiled result.
        """

        argtypes = [
            dpex_kernel_target.typing_context.resolve_argument_type(arg)
            for arg in args
        ]

        # Generate key used for cache lookup
        stripped_argtypes = strip_usm_metadata(argtypes)
        codegen_magic_tuple = (
            dpex_kernel_target.target_context.codegen().magic_tuple()
        )
        key = build_key(stripped_argtypes, codegen_magic_tuple, self._func_hash)

        cres = self._cache.get(key)
        if cres is None:
            self._cache_hits += 1
            cres = compile_with_dpex(
                pyfunc=self._pyfunc,
                pyfunc_name=self._pyfunc.__name__,
                return_type=None,
                target_context=dpex_kernel_target.target_context,
                typing_context=dpex_kernel_target.typing_context,
                args=args,
                is_kernel=False,
                debug=self._debug,
            )
            func = cres.library.get_function(cres.fndesc.llvm_func_name)
            cres.target_context.set_spir_func_calling_conv(func)
            libs = [cres.library]

            cres.target_context.insert_user_function(self, cres.fndesc, libs)
            self._cache.put(key, cres)
        return cres.signature


def compile_func(pyfunc, signature, debug=False):
    """Compiles a specialized `numba_dpex.func`

    Compiles a specialized `numba_dpex.func` decorated function to native binary
    library function and returns the library wrapped inside a
    `numba_dpex.core.kernel_interface.func.DpexFunction` object.

    Args:
        pyfunc (`function`): A python function to be compiled.
        signature (`list`): A list of `numba.core.typing.templates.Signature`'s
        debug (`bool`, optional): Debug options. Defaults to `False`.

    Returns:
        `numba_dpex.core.kernel_interface.func.DpexFunction`: A `DpexFunction`
         object
    """

    devfn = DpexFunction(pyfunc, debug=debug)

    cres = []
    for sig in signature:
        arg_types, return_types = sigutils.normalize_signature(sig)
        c = devfn.compile(arg_types, return_types)
        cres.append(c)

    class _function_template(ConcreteTemplate):
        unsafe_casting = False
        exact_match_required = True
        key = devfn
        cases = [c.signature for c in cres]

    cres[0].typing_context.insert_user_function(devfn, _function_template)

    for c in cres:
        c.target_context.insert_user_function(devfn, c.fndesc, [c.library])

    return devfn


def compile_func_template(pyfunc, debug=False, enable_cache=True):
    """Converts a `numba_dpex.func` function to an `AbstractTemplate`

    Converts a `numba_dpex.func` decorated function to a Numba
    `AbstractTemplate` and returns the object wrapped inside a
    `numba_dpex.core.kernel_interface.func.DpexFunctionTemplate`
    object.

    A `DpexFunctionTemplate` object is an abstract representation for
    a native function with `spir_func` calling convention that is to be
    JIT compiled once the argument types are resolved.

    Args:
        pyfunc (`function`): A python function to be compiled.
        debug (`bool`, optional): Debug options. Defaults to `False`.

    Raises:
        `AssertionError`: Raised if keyword arguments are supplied in
            the inner generic function.

    Returns:
        `numba_dpex.core.kernel_interface.func.DpexFunctionTemplate`:
            A `DpexFunctionTemplate` object.
    """

    dft = DpexFunctionTemplate(pyfunc, debug=debug, enable_cache=enable_cache)

    class _function_template(AbstractTemplate):
        unsafe_casting = False
        exact_match_required = True
        key = dft

        def generic(self, args, kws):
            if kws:
                raise AssertionError("No keyword arguments allowed.")
            return dft.compile(args)

    dpex_kernel_target.typing_context.insert_user_function(
        dft, _function_template
    )

    return dft

1	# SPDX-FileCopyrightText: 2022 - 2024 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
UNCOV 5	from numba.core import sigutils, types	×
UNCOV 6	from numba.core.typing.templates import AbstractTemplate, ConcreteTemplate	×
7
UNCOV 8	from numba_dpex.core import config	×
UNCOV 9	from numba_dpex.core.caching import LRUCache, NullCache	×
UNCOV 10	from numba_dpex.core.compiler import compile_with_dpex	×
UNCOV 11	from numba_dpex.core.descriptor import dpex_kernel_target	×
UNCOV 12	from numba_dpex.core.utils import (	×
13	build_key,
14	create_func_hash,
15	strip_usm_metadata,
16	)
17
18
UNCOV 19	class DpexFunction(object):	×
20	"""Class to materialize dpex function
21
22	Helper class to eager compile a specialized `numba_dpex.func`
23	decorated Python function into a LLVM function with `spir_func`
24	calling convention.
25
26	A specialized `numba_dpex.func` decorated Python function is one
27	where the user has specified a signature or a list of signatures
28	for the function. The function gets compiled as soon as the Python
29	program is loaded, i.e., eagerly, instead of JIT compilation once
30	the function is invoked.
31	"""
32
UNCOV 33	def __init__(self, pyfunc, debug=False):	×
34	"""Constructor for `DpexFunction`
35
36	Args:
37	pyfunc (`function`): A python function to be compiled.
38	debug (`bool`, optional): Debug option for compilation.
39	Defaults to `False`.
40	"""
41	self._pyfunc = pyfunc	×
42	self._debug = debug	×
43
UNCOV 44	def compile(self, arg_types, return_types):	×
45	"""The actual compilation function.
46
47	Args:
48	arg_types (`tuple`): Function argument types in a tuple.
49	return_types (`numba.core.types.scalars.Integer`):
50	An integer value to specify the return type.
51
52	Returns:
53	`numba.core.compiler.CompileResult`: The compiled result
54	"""
55
56	cres = compile_with_dpex(	×
57	pyfunc=self._pyfunc,
58	pyfunc_name=self._pyfunc.__name__,
59	return_type=return_types,
60	target_context=dpex_kernel_target.target_context,
61	typing_context=dpex_kernel_target.typing_context,
62	args=arg_types,
63	is_kernel=False,
64	debug=self._debug,
65	)
66	func = cres.library.get_function(cres.fndesc.llvm_func_name)	×
67	cres.target_context.set_spir_func_calling_conv(func)	×
68
69	return cres	×
70
71
UNCOV 72	class DpexFunctionTemplate(object):	×
73	"""Helper class to compile an unspecialized `numba_dpex.func`
74
75	A helper class to JIT compile an unspecialized `numba_dpex.func`
76	decorated Python function into an LLVM function with `spir_func`
77	calling convention.
78	"""
79
UNCOV 80	def __init__(self, pyfunc, debug=False, enable_cache=True):	×
81	"""Constructor for `DpexFunctionTemplate`
82
83	Args:
84	pyfunc (function): A python function to be compiled.
85	debug (bool, optional): Debug option for compilation.
86	Defaults to `False`.
87	enable_cache (bool, optional): Flag to turn on/off caching.
88	Defaults to `True`.
89	"""
UNCOV 90	self._pyfunc = pyfunc	×
UNCOV 91	self._debug = debug	×
UNCOV 92	self._enable_cache = enable_cache	×
93
UNCOV 94	self._func_hash = create_func_hash(pyfunc)	×
95
UNCOV 96	if not config.ENABLE_CACHE:	×
97	self._cache = NullCache()	×
UNCOV 98	elif self._enable_cache:	×
UNCOV 99	self._cache = LRUCache(	×
100	name="DpexFunctionTemplateCache",
101	capacity=config.CACHE_SIZE,
102	pyfunc=self._pyfunc,
103	)
104	else:
105	self._cache = NullCache()	×
UNCOV 106	self._cache_hits = 0	×
107
UNCOV 108	@property	×
UNCOV 109	def cache(self):	×
110	"""Cache accessor"""
111	return self._cache	×
112
UNCOV 113	@property	×
UNCOV 114	def cache_hits(self):	×
115	"""Cache hit count accessor"""
116	return self._cache_hits	×
117
UNCOV 118	def compile(self, args):	×
119	"""Compile a `numba_dpex.func` decorated function
120
121	Compile a `numba_dpex.func` decorated Python function with the
122	given argument types. Each signature is compiled once by caching
123	the compiled function inside this object.
124
125	Args:
126	args (`tuple`): Function argument types in a tuple.
127
128	Returns:
129	`numba.core.typing.templates.Signature`: Signature of the
130	compiled result.
131	"""
132
UNCOV 133	argtypes = [	×
134	dpex_kernel_target.typing_context.resolve_argument_type(arg)
135	for arg in args
136	]
137
138	# Generate key used for cache lookup
UNCOV 139	stripped_argtypes = strip_usm_metadata(argtypes)	×
UNCOV 140	codegen_magic_tuple = (	×
141	dpex_kernel_target.target_context.codegen().magic_tuple()
142	)
UNCOV 143	key = build_key(stripped_argtypes, codegen_magic_tuple, self._func_hash)	×
144
UNCOV 145	cres = self._cache.get(key)	×
UNCOV 146	if cres is None:	×
UNCOV 147	self._cache_hits += 1	×
UNCOV 148	cres = compile_with_dpex(	×
149	pyfunc=self._pyfunc,
150	pyfunc_name=self._pyfunc.__name__,
151	return_type=None,
152	target_context=dpex_kernel_target.target_context,
153	typing_context=dpex_kernel_target.typing_context,
154	args=args,
155	is_kernel=False,
156	debug=self._debug,
157	)
UNCOV 158	func = cres.library.get_function(cres.fndesc.llvm_func_name)	×
UNCOV 159	cres.target_context.set_spir_func_calling_conv(func)	×
UNCOV 160	libs = [cres.library]	×
161
UNCOV 162	cres.target_context.insert_user_function(self, cres.fndesc, libs)	×
UNCOV 163	self._cache.put(key, cres)	×
UNCOV 164	return cres.signature	×
165
166
UNCOV 167	def compile_func(pyfunc, signature, debug=False):	×
168	"""Compiles a specialized `numba_dpex.func`
169
170	Compiles a specialized `numba_dpex.func` decorated function to native binary
171	library function and returns the library wrapped inside a
172	`numba_dpex.core.kernel_interface.func.DpexFunction` object.
173
174	Args:
175	pyfunc (`function`): A python function to be compiled.
176	signature (`list`): A list of `numba.core.typing.templates.Signature`'s
177	debug (`bool`, optional): Debug options. Defaults to `False`.
178
179	Returns:
180	`numba_dpex.core.kernel_interface.func.DpexFunction`: A `DpexFunction`
181	object
182	"""
183
184	devfn = DpexFunction(pyfunc, debug=debug)	×
185
186	cres = []	×
187	for sig in signature:	×
188	arg_types, return_types = sigutils.normalize_signature(sig)	×
189	c = devfn.compile(arg_types, return_types)	×
190	cres.append(c)	×
191
192	class _function_template(ConcreteTemplate):	×
193	unsafe_casting = False	×
194	exact_match_required = True	×
195	key = devfn	×
196	cases = [c.signature for c in cres]	×
197
198	cres[0].typing_context.insert_user_function(devfn, _function_template)	×
199
200	for c in cres:	×
201	c.target_context.insert_user_function(devfn, c.fndesc, [c.library])	×
202
203	return devfn	×
204
205
UNCOV 206	def compile_func_template(pyfunc, debug=False, enable_cache=True):	×
207	"""Converts a `numba_dpex.func` function to an `AbstractTemplate`
208
209	Converts a `numba_dpex.func` decorated function to a Numba
210	`AbstractTemplate` and returns the object wrapped inside a
211	`numba_dpex.core.kernel_interface.func.DpexFunctionTemplate`
212	object.
213
214	A `DpexFunctionTemplate` object is an abstract representation for
215	a native function with `spir_func` calling convention that is to be
216	JIT compiled once the argument types are resolved.
217
218	Args:
219	pyfunc (`function`): A python function to be compiled.
220	debug (`bool`, optional): Debug options. Defaults to `False`.
221
222	Raises:
223	`AssertionError`: Raised if keyword arguments are supplied in
224	the inner generic function.
225
226	Returns:
227	`numba_dpex.core.kernel_interface.func.DpexFunctionTemplate`:
228	A `DpexFunctionTemplate` object.
229	"""
230
UNCOV 231	dft = DpexFunctionTemplate(pyfunc, debug=debug, enable_cache=enable_cache)	×
232
UNCOV 233	class _function_template(AbstractTemplate):	×
UNCOV 234	unsafe_casting = False	×
UNCOV 235	exact_match_required = True	×
UNCOV 236	key = dft	×
237
UNCOV 238	def generic(self, args, kws):	×
UNCOV 239	if kws:	×
240	raise AssertionError("No keyword arguments allowed.")	×
UNCOV 241	return dft.compile(args)	×
242
UNCOV 243	dpex_kernel_target.typing_context.insert_user_function(	×
244	dft, _function_template
245	)
246
UNCOV 247	return dft	×

IntelPython / numba-dpex / 8386358402

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