8289101957

Committed 14 Mar 2024 11:53PM UTC coverage: 82.62% (-0.02%) from 82.64%

Build # 8289101957

Build Type

Pull #1377

github

Committed by

web-flow

Commit Message

Merge 01318388a into 283186249

Pull Request Pull Request #1377: Fix/private array

Run Details

1567 of 2165 branches covered (72.38%)

Branch coverage included in aggregate %.

47 of 54 new or added lines in 7 files covered. (87.04%)

1 existing line in 1 file now uncovered.

6462 of 7553 relevant lines covered (85.56%)

0.86 hits per line

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

80.86

/numba_dpex/kernel_api_impl/spirv/dispatcher.py

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

"""Implements a new numba dispatcher class and a compiler class to compile and
call numba_dpex.kernel decorated function.
"""
import hashlib
from collections import namedtuple
from contextlib import ExitStack
from typing import Tuple

import numba.core.event as ev
from llvmlite.binding.value import ValueRef
from numba.core import errors, sigutils, types
from numba.core.compiler import CompileResult, Flags
from numba.core.compiler_lock import global_compiler_lock
from numba.core.dispatcher import Dispatcher, _FunctionCompiler
from numba.core.funcdesc import PythonFunctionDescriptor
from numba.core.target_extension import (
    dispatcher_registry,
    target_override,
    target_registry,
)
from numba.core.types import Array as NpArrayType
from numba.core.types import void
from numba.core.typing.typeof import Purpose, typeof

from numba_dpex import config, numba_sem_version
from numba_dpex.core.exceptions import (
    ExecutionQueueInferenceError,
    InvalidKernelSpecializationError,
    KernelHasReturnValueError,
    UnsupportedKernelArgumentError,
)
from numba_dpex.core.pipelines import kernel_compiler
from numba_dpex.core.types import USMNdArray
from numba_dpex.core.utils import kernel_launcher as kl
from numba_dpex.experimental.target import (
    DPEX_KERNEL_EXP_TARGET_NAME,
    dpex_exp_kernel_target,
)
from numba_dpex.kernel_api_impl.spirv import spirv_generator
from numba_dpex.kernel_api_impl.spirv.codegen import SPIRVCodeLibrary
from numba_dpex.kernel_api_impl.spirv.target import (
    CompilationMode,
    SPIRVTargetContext,
)

_SPIRVKernelCompileResult = namedtuple(
    "_KernelCompileResult", CompileResult._fields + ("kernel_device_ir_module",)
)


class _SPIRVKernelCompiler(_FunctionCompiler):
    """A special compiler class used to compile numba_dpex.kernel decorated
    functions.
    """

    def check_queue_equivalence_of_args(
        self, py_func_name: str, args: [types.Type, ...]
    ):
        """Evaluates if all USMNdArray arguments passed to a kernel function
        has the same DpctlSyclQueue type.

        Args:
            py_func_name (str): Name of the kernel that is being evaluated
            args (types.Type, ...]): List of numba inferred types for each
            argument passed to the kernel

        Raises:
            ExecutionQueueInferenceError: If all USMNdArray were not allocated
            on the same dpctl.SyclQueue
            ExecutionQueueInferenceError: If there were not USMNdArray
            arguments passed to the kernel.
        """
        common_queue = None

        for arg in args:
            if isinstance(arg, USMNdArray):
                if common_queue is None:
                    common_queue = arg.queue
                elif common_queue != arg.queue:
                    raise ExecutionQueueInferenceError(
                        kernel_name=py_func_name, usmarray_argnum_list=[]
                    )

        if common_queue is None:
            raise ExecutionQueueInferenceError(
                kernel_name=py_func_name, usmarray_argnum_list=None
            )

    def check_sig_types(self, py_func_name: str, argtypes, return_type):
        """Checks signature type and raises error if wrong signature was
        provided.

        Args:
            sig: The signature on which the kernel is to be specialized.

        Raises:
            KernelHasReturnValueError: non void return type.
            InvalidKernelSpecializationError: unsupported arguments where
                provided.
        """

        if return_type is None:
            return_type = void
        sig = return_type(*argtypes)

        # Check if signature has a non-void return type
        if return_type and return_type != void:
            raise KernelHasReturnValueError(
                kernel_name=None,
                return_type=return_type,
                sig=sig,
            )

        # USMNdarray check
        usmarray_argnums = []
        usmndarray_argtypes = []
        unsupported_argnum_list = []

        for i, argtype in enumerate(argtypes):
            # FIXME: Add checks for other types of unsupported kernel args, e.g.
            # complex.

            # Check if a non-USMNdArray Array type is passed to the kernel
            if isinstance(argtype, NpArrayType) and not isinstance(
                argtype, USMNdArray
            ):
                unsupported_argnum_list.append(i)
            elif isinstance(argtype, USMNdArray):
                usmarray_argnums.append(i)
                usmndarray_argtypes.append(argtype)

        if unsupported_argnum_list:
            raise InvalidKernelSpecializationError(
                kernel_name=py_func_name,
                invalid_sig=sig,
                unsupported_argnum_list=unsupported_argnum_list,
            )

    def check_arguments(self, py_func_name: str, args: [types.Type, ...]):
        """Checks arguments and queue of the input arguments.

        Raises:
            KernelHasReturnValueError: non void return type.
            InvalidKernelSpecializationError: unsupported arguments where
                provided.
            ExecutionQueueInferenceError: If all USMNdArray were not allocated
                on the same dpctl.SyclQueue
            ExecutionQueueInferenceError: If there were not USMNdArray
                arguments passed to the kernel.
        """
        self.check_sig_types(py_func_name, args, None)
        self.check_queue_equivalence_of_args(py_func_name, args)

    def _compile_to_spirv(
        self,
        kernel_library: SPIRVCodeLibrary,
        kernel_fndesc: PythonFunctionDescriptor,
        kernel_targetctx: SPIRVTargetContext,
    ):
        kernel_func: ValueRef = kernel_library.get_function(
            kernel_fndesc.llvm_func_name
        )

        # Create a spir_kernel wrapper function
        kernel_fn = kernel_targetctx.prepare_spir_kernel(
            kernel_func, kernel_fndesc.argtypes
        )
        # Get the compiler flags that were passed through the target descriptor
        flags = Flags()
        self.targetdescr.options.parse_as_flags(flags, self.targetoptions)

        # If the inline_threshold option was set then set the property in the
        # kernel_library to force inlining ``overload`` calls into a kernel.
        inline_threshold = flags.inline_threshold  # pylint: disable=E1101
        kernel_library.inline_threshold = inline_threshold

        # Call finalize on the LLVM module. Finalization will result in
        # all linking libraries getting linked together and final optimization
        # including inlining of functions if an inlining level is specified.
        kernel_library.finalize()

        if config.DUMP_KERNEL_LLVM:
            self._dump_kernel(kernel_fndesc, kernel_library)
        # Compiled the LLVM IR to SPIR-V
        kernel_spirv_module = spirv_generator.llvm_to_spirv(
            kernel_targetctx,
            kernel_library.final_module,
            kernel_library.final_module.as_bitcode(),
        )
        return kl.SPIRVKernelModule(
            kernel_name=kernel_fn.name, kernel_bitcode=kernel_spirv_module
        )

    def compile(self, args, return_type) -> _SPIRVKernelCompileResult:
        status, kcres = self._compile_cached(args, return_type)
        if status:
            return kcres

        raise kcres

    def _compile_cached(
        self, args, return_type: types.Type
    ) -> Tuple[bool, _SPIRVKernelCompileResult]:
        """Compiles the kernel function to bitcode and generates a host-callable
        wrapper to submit the kernel to a SYCL queue.

        The LLVM IR generated for the kernel function is available in the
        CompileResult objected returned by
        numba_dpex.core.pipeline.kernel_compiler.KernelCompiler.

        Once the kernel decorated function is compiled down to LLVM IR, the
        following steps are performed:

            a) compile the IR into SPIR-V kernel
            b) generate a host callable wrapper function that will create a
               sycl::kernel_bundle from the SPIR-V and then submits the
               kernel_bundle to a sycl::queue
            c) create a cpython_wrapper_function for the host callable wrapper
               function.
            d) create a cfunc_wrapper_function to make the host callable wrapper
               function callable inside another JIT-compiled function.

        Args:
            args (tuple(types.Type)): A tuple of numba.core.Type instances each
            representing the numba-inferred type of a kernel argument.

            return_type (types.Type): The numba-inferred type of the returned
            value from the kernel. Should always be types.NoneType.

        Returns:
            CompileResult: A CompileResult object storing the LLVM library for
            the host-callable wrapper function.
        """
        key = tuple(args), return_type
        try:
            return False, self._failed_cache[key]
        except KeyError:
            pass

        try:
            with target_override(DPEX_KERNEL_EXP_TARGET_NAME):
                cres: CompileResult = self._compile_core(args, return_type)

            if (
                self.targetoptions["_compilation_mode"]
                == CompilationMode.KERNEL
            ):
                kernel_device_ir_module: kl.SPIRVKernelModule = (
                    self._compile_to_spirv(
                        cres.library, cres.fndesc, cres.target_context
                    )
                )
            else:
                kernel_device_ir_module = None

            kcres_attrs = []

            for cres_field in cres._fields:
                cres_attr = getattr(cres, cres_field)
                if cres_field == "entry_point":
                    if cres_attr is not None:
                        raise AssertionError(
                            "Compiled kernel and device_func should be "
                            "compiled with compile_cfunc option turned off"
                        )
                    cres_attr = cres.fndesc.qualname
                kcres_attrs.append(cres_attr)

            kcres_attrs.append(kernel_device_ir_module)

        except errors.TypingError as err:
            self._failed_cache[key] = err
            return False, err

        return True, _SPIRVKernelCompileResult(*kcres_attrs)

    def _dump_kernel(self, fndesc, library):
        """Dump kernel into file."""
        name = fndesc.llvm_func_name
        if len(name) > 200:
            sha256 = hashlib.sha256(name.encode("utf-8")).hexdigest()
            name = name[:150] + "_" + sha256

        with open(
            name + ".ll",
            "w",
            encoding="UTF-8",
        ) as fptr:
            fptr.write(str(library.final_module))


class SPIRVKernelDispatcher(Dispatcher):
    """Dispatcher class designed to compile kernel decorated functions. The
    dispatcher inherits the Numba Dispatcher class, but has a different
    compilation strategy. Instead of compiling a kernel decorated function to
    an executable binary, the dispatcher compiles it to SPIR-V and then caches
    that SPIR-V bitcode.

    """

    targetdescr = dpex_exp_kernel_target
    _fold_args = False

    def __init__(
        self,
        pyfunc,
        local_vars_to_numba_types=None,
        targetoptions=None,
        pipeline_class=kernel_compiler.KernelCompiler,
    ):
        if targetoptions is None:
            targetoptions = {}

        if local_vars_to_numba_types is None:
            local_vars_to_numba_types = {}

        targetoptions["nopython"] = True
        targetoptions["experimental"] = True

        self._kernel_name = pyfunc.__name__

        if numba_sem_version < (0, 59, 0):
            # pylint: disable=unexpected-keyword-arg
            super().__init__(
                py_func=pyfunc,
                locals=local_vars_to_numba_types,
                impl_kind="direct",
                targetoptions=targetoptions,
                pipeline_class=pipeline_class,
            )
        else:
            super().__init__(
                py_func=pyfunc,
                locals=local_vars_to_numba_types,
                targetoptions=targetoptions,
                pipeline_class=pipeline_class,
            )
        self._compiler = _SPIRVKernelCompiler(
            pyfunc,
            self.targetdescr,
            targetoptions,
            local_vars_to_numba_types,
            pipeline_class,
        )

    def typeof_pyval(self, val):
        """
        Resolve the Numba type of Python value *val*.
        This is called from numba._dispatcher as a fallback if the native code
        cannot decide the type.
        """
        # Not going through the resolve_argument_type() indirection
        # can save a couple µs.
        try:
            typ = typeof(val, Purpose.argument)
            if isinstance(typ, types.Array) and not isinstance(typ, USMNdArray):
                raise UnsupportedKernelArgumentError(
                    type=str(type(val)), value=val
                )
        except ValueError:
            typ = types.pyobject
        else:
            if typ is None:
                typ = types.pyobject
        self._types_active_call.append(typ)
        return typ

    def add_overload(self, cres):
        args = tuple(cres.signature.args)
        self.overloads[args] = cres

    def compile(self, sig) -> any:
        disp = self._get_dispatcher_for_current_target()
        if disp is not self:
            return disp.compile(sig)

        with ExitStack() as scope:
            cres = None

            def cb_compiler(dur):
                if cres is not None:
                    self._callback_add_compiler_timer(dur, cres)

            def cb_llvm(dur):
                if cres is not None:
                    self._callback_add_llvm_timer(dur, cres)

            scope.enter_context(
                ev.install_timer("numba:compiler_lock", cb_compiler)
            )
            scope.enter_context(ev.install_timer("numba:llvm_lock", cb_llvm))
            scope.enter_context(global_compiler_lock)

            if not self._can_compile:
                raise RuntimeError("compilation disabled")
            # Use counter to track recursion compilation depth
            with self._compiling_counter:
                args, return_type = sigutils.normalize_signature(sig)

                if (
                    self.targetoptions["_compilation_mode"]
                    == CompilationMode.KERNEL
                ):
                    # Compute follows data based queue equivalence is only
                    # evaluated for kernel functions whose arguments are
                    # supposed to be arrays. For device_func decorated
                    # functions, the arguments can be scalar and we skip queue
                    # equivalence check.
                    try:
                        self._compiler.check_arguments(self._kernel_name, args)
                    except ExecutionQueueInferenceError as eqie:
                        raise eqie

                    # A function being compiled in the KERNEL compilation mode
                    # cannot have a non-void return value
                    if return_type and return_type != void:
                        raise KernelHasReturnValueError(
                            kernel_name=None, return_type=return_type, sig=sig
                        )

                # Don't recompile if signature already exists
                existing = self.overloads.get(tuple(args))
                if existing is not None:
                    return existing.entry_point

                # TODO: Enable caching
                # Add code to enable on disk caching of a binary spirv kernel.
                # Refer: https://github.com/IntelPython/numba-dpex/issues/1197
                self._cache_misses[sig] += 1
                with ev.trigger_event(
                    "numba_dpex:compile",
                    data={
                        "dispatcher": self,
                        "args": args,
                        "return_type": return_type,
                    },
                ):
                    try:
                        compiler: _SPIRVKernelCompiler = self._compiler
                        kcres: _SPIRVKernelCompileResult = compiler.compile(
                            args, return_type
                        )
                    except errors.ForceLiteralArg as err:

                        def folded(args, kws):
                            return self._compiler.fold_argument_types(
                                args, kws
                            )[1]

                        raise err.bind_fold_arguments(folded)
                    self.add_overload(kcres)

                    kcres.target_context.insert_user_function(
                        kcres.entry_point, kcres.fndesc, [kcres.library]
                    )

                # TODO: enable caching of kernel_module
                # https://github.com/IntelPython/numba-dpex/issues/1197

                return kcres.entry_point

    def __getitem__(self, args):
        """Square-bracket notation for configuring launch arguments is not
        supported.
        """

        raise NotImplementedError

    def __call__(self, *args, **kw_args):
        """Functor to launch a kernel."""

        raise NotImplementedError


_dpex_target = target_registry[DPEX_KERNEL_EXP_TARGET_NAME]
dispatcher_registry[_dpex_target] = SPIRVKernelDispatcher

1	# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	"""Implements a new numba dispatcher class and a compiler class to compile and	1✔
6	call numba_dpex.kernel decorated function.
7	"""
8	import hashlib	1✔
9	from collections import namedtuple	1✔
10	from contextlib import ExitStack	1✔
11	from typing import Tuple	1✔
12
13	import numba.core.event as ev	1✔
14	from llvmlite.binding.value import ValueRef	1✔
15	from numba.core import errors, sigutils, types	1✔
16	from numba.core.compiler import CompileResult, Flags	1✔
17	from numba.core.compiler_lock import global_compiler_lock	1✔
18	from numba.core.dispatcher import Dispatcher, _FunctionCompiler	1✔
19	from numba.core.funcdesc import PythonFunctionDescriptor	1✔
20	from numba.core.target_extension import (	1✔
21	dispatcher_registry,
22	target_override,
23	target_registry,
24	)
25	from numba.core.types import Array as NpArrayType	1✔
26	from numba.core.types import void	1✔
27	from numba.core.typing.typeof import Purpose, typeof	1✔
28
29	from numba_dpex import config, numba_sem_version	1✔
30	from numba_dpex.core.exceptions import (	1✔
31	ExecutionQueueInferenceError,
32	InvalidKernelSpecializationError,
33	KernelHasReturnValueError,
34	UnsupportedKernelArgumentError,
35	)
36	from numba_dpex.core.pipelines import kernel_compiler	1✔
37	from numba_dpex.core.types import USMNdArray	1✔
38	from numba_dpex.core.utils import kernel_launcher as kl	1✔
39	from numba_dpex.experimental.target import (	1✔
40	DPEX_KERNEL_EXP_TARGET_NAME,
41	dpex_exp_kernel_target,
42	)
43	from numba_dpex.kernel_api_impl.spirv import spirv_generator	1✔
44	from numba_dpex.kernel_api_impl.spirv.codegen import SPIRVCodeLibrary	1✔
45	from numba_dpex.kernel_api_impl.spirv.target import (	1✔
46	CompilationMode,
47	SPIRVTargetContext,
48	)
49
50	_SPIRVKernelCompileResult = namedtuple(	1✔
51	"_KernelCompileResult", CompileResult._fields + ("kernel_device_ir_module",)
52	)
53
54
55	class _SPIRVKernelCompiler(_FunctionCompiler):	1✔
56	"""A special compiler class used to compile numba_dpex.kernel decorated
57	functions.
58	"""
59
60	def check_queue_equivalence_of_args(	1✔
61	self, py_func_name: str, args: [types.Type, ...]
62	):
63	"""Evaluates if all USMNdArray arguments passed to a kernel function
64	has the same DpctlSyclQueue type.
65
66	Args:
67	py_func_name (str): Name of the kernel that is being evaluated
68	args (types.Type, ...]): List of numba inferred types for each
69	argument passed to the kernel
70
71	Raises:
72	ExecutionQueueInferenceError: If all USMNdArray were not allocated
73	on the same dpctl.SyclQueue
74	ExecutionQueueInferenceError: If there were not USMNdArray
75	arguments passed to the kernel.
76	"""
77	common_queue = None	1✔
78
79	for arg in args:	1✔
80	if isinstance(arg, USMNdArray):	1✔
81	if common_queue is None:	1✔
82	common_queue = arg.queue	1✔
83	elif common_queue != arg.queue:	1✔
84	raise ExecutionQueueInferenceError(	1✔
85	kernel_name=py_func_name, usmarray_argnum_list=[]
86	)
87
88	if common_queue is None:	1✔
89	raise ExecutionQueueInferenceError(	1✔
90	kernel_name=py_func_name, usmarray_argnum_list=None
91	)
92
93	def check_sig_types(self, py_func_name: str, argtypes, return_type):	1✔
94	"""Checks signature type and raises error if wrong signature was
95	provided.
96
97	Args:
98	sig: The signature on which the kernel is to be specialized.
99
100	Raises:
101	KernelHasReturnValueError: non void return type.
102	InvalidKernelSpecializationError: unsupported arguments where
103	provided.
104	"""
105
106	if return_type is None:	1!
107	return_type = void	1✔
108	sig = return_type(*argtypes)	1✔
109
110	# Check if signature has a non-void return type
111	if return_type and return_type != void:	1!
112	raise KernelHasReturnValueError(	×
113	kernel_name=None,
114	return_type=return_type,
115	sig=sig,
116	)
117
118	# USMNdarray check
119	usmarray_argnums = []	1✔
120	usmndarray_argtypes = []	1✔
121	unsupported_argnum_list = []	1✔
122
123	for i, argtype in enumerate(argtypes):	1✔
124	# FIXME: Add checks for other types of unsupported kernel args, e.g.
125	# complex.
126
127	# Check if a non-USMNdArray Array type is passed to the kernel
128	if isinstance(argtype, NpArrayType) and not isinstance(	1✔
129	argtype, USMNdArray
130	):
131	unsupported_argnum_list.append(i)	1✔
132	elif isinstance(argtype, USMNdArray):	1✔
133	usmarray_argnums.append(i)	1✔
134	usmndarray_argtypes.append(argtype)	1✔
135
136	if unsupported_argnum_list:	1✔
137	raise InvalidKernelSpecializationError(	1✔
138	kernel_name=py_func_name,
139	invalid_sig=sig,
140	unsupported_argnum_list=unsupported_argnum_list,
141	)
142
143	def check_arguments(self, py_func_name: str, args: [types.Type, ...]):	1✔
144	"""Checks arguments and queue of the input arguments.
145
146	Raises:
147	KernelHasReturnValueError: non void return type.
148	InvalidKernelSpecializationError: unsupported arguments where
149	provided.
150	ExecutionQueueInferenceError: If all USMNdArray were not allocated
151	on the same dpctl.SyclQueue
152	ExecutionQueueInferenceError: If there were not USMNdArray
153	arguments passed to the kernel.
154	"""
155	self.check_sig_types(py_func_name, args, None)	1✔
156	self.check_queue_equivalence_of_args(py_func_name, args)	1✔
157
158	def _compile_to_spirv(	1✔
159	self,
160	kernel_library: SPIRVCodeLibrary,
161	kernel_fndesc: PythonFunctionDescriptor,
162	kernel_targetctx: SPIRVTargetContext,
163	):
164	kernel_func: ValueRef = kernel_library.get_function(	1✔
165	kernel_fndesc.llvm_func_name
166	)
167
168	# Create a spir_kernel wrapper function
169	kernel_fn = kernel_targetctx.prepare_spir_kernel(	1✔
170	kernel_func, kernel_fndesc.argtypes
171	)
172	# Get the compiler flags that were passed through the target descriptor
173	flags = Flags()	1✔
174	self.targetdescr.options.parse_as_flags(flags, self.targetoptions)	1✔
175
176	# If the inline_threshold option was set then set the property in the
177	# kernel_library to force inlining ``overload`` calls into a kernel.
178	inline_threshold = flags.inline_threshold # pylint: disable=E1101	1✔
179	kernel_library.inline_threshold = inline_threshold	1✔
180
181	# Call finalize on the LLVM module. Finalization will result in
182	# all linking libraries getting linked together and final optimization
183	# including inlining of functions if an inlining level is specified.
184	kernel_library.finalize()	1✔
185
186	if config.DUMP_KERNEL_LLVM:	1!
NEW 187	self._dump_kernel(kernel_fndesc, kernel_library)	×
188	# Compiled the LLVM IR to SPIR-V
189	kernel_spirv_module = spirv_generator.llvm_to_spirv(	1✔
190	kernel_targetctx,
191	kernel_library.final_module,
192	kernel_library.final_module.as_bitcode(),
193	)
194	return kl.SPIRVKernelModule(	1✔
195	kernel_name=kernel_fn.name, kernel_bitcode=kernel_spirv_module
196	)
197
198	def compile(self, args, return_type) -> _SPIRVKernelCompileResult:	1✔
199	status, kcres = self._compile_cached(args, return_type)	1✔
200	if status:	1✔
201	return kcres	1✔
202
203	raise kcres	1✔
204
205	def _compile_cached(	1✔
206	self, args, return_type: types.Type
207	) -> Tuple[bool, _SPIRVKernelCompileResult]:
208	"""Compiles the kernel function to bitcode and generates a host-callable
209	wrapper to submit the kernel to a SYCL queue.
210
211	The LLVM IR generated for the kernel function is available in the
212	CompileResult objected returned by
213	numba_dpex.core.pipeline.kernel_compiler.KernelCompiler.
214
215	Once the kernel decorated function is compiled down to LLVM IR, the
216	following steps are performed:
217
218	a) compile the IR into SPIR-V kernel
219	b) generate a host callable wrapper function that will create a
220	sycl::kernel_bundle from the SPIR-V and then submits the
221	kernel_bundle to a sycl::queue
222	c) create a cpython_wrapper_function for the host callable wrapper
223	function.
224	d) create a cfunc_wrapper_function to make the host callable wrapper
225	function callable inside another JIT-compiled function.
226
227	Args:
228	args (tuple(types.Type)): A tuple of numba.core.Type instances each
229	representing the numba-inferred type of a kernel argument.
230
231	return_type (types.Type): The numba-inferred type of the returned
232	value from the kernel. Should always be types.NoneType.
233
234	Returns:
235	CompileResult: A CompileResult object storing the LLVM library for
236	the host-callable wrapper function.
237	"""
238	key = tuple(args), return_type	1✔
239	try:	1✔
240	return False, self._failed_cache[key]	1✔
241	except KeyError:	1✔
242	pass	1✔
243
244	try:	1✔
245	with target_override(DPEX_KERNEL_EXP_TARGET_NAME):	1✔
246	cres: CompileResult = self._compile_core(args, return_type)	1✔
247
248	if (	1✔
249	self.targetoptions["_compilation_mode"]
250	== CompilationMode.KERNEL
251	):
252	kernel_device_ir_module: kl.SPIRVKernelModule = (	1✔
253	self._compile_to_spirv(
254	cres.library, cres.fndesc, cres.target_context
255	)
256	)
257	else:
258	kernel_device_ir_module = None	1✔
259
260	kcres_attrs = []	1✔
261
262	for cres_field in cres._fields:	1✔
263	cres_attr = getattr(cres, cres_field)	1✔
264	if cres_field == "entry_point":	1✔
265	if cres_attr is not None:	1!
266	raise AssertionError(	×
267	"Compiled kernel and device_func should be "
268	"compiled with compile_cfunc option turned off"
269	)
270	cres_attr = cres.fndesc.qualname	1✔
271	kcres_attrs.append(cres_attr)	1✔
272
273	kcres_attrs.append(kernel_device_ir_module)	1✔
274
275	except errors.TypingError as err:	1✔
276	self._failed_cache[key] = err	1✔
277	return False, err	1✔
278
279	return True, _SPIRVKernelCompileResult(*kcres_attrs)	1✔
280
281	def _dump_kernel(self, fndesc, library):	1✔
282	"""Dump kernel into file."""
NEW 283	name = fndesc.llvm_func_name	×
NEW 284	if len(name) > 200:	×
NEW 285	sha256 = hashlib.sha256(name.encode("utf-8")).hexdigest()	×
NEW 286	name = name[:150] + "_" + sha256	×
287
NEW 288	with open(	×
289	name + ".ll",
290	"w",
291	encoding="UTF-8",
292	) as fptr:
NEW 293	fptr.write(str(library.final_module))	×
294
295
296	class SPIRVKernelDispatcher(Dispatcher):	1✔
297	"""Dispatcher class designed to compile kernel decorated functions. The
298	dispatcher inherits the Numba Dispatcher class, but has a different
299	compilation strategy. Instead of compiling a kernel decorated function to
300	an executable binary, the dispatcher compiles it to SPIR-V and then caches
301	that SPIR-V bitcode.
302
303	"""
304
305	targetdescr = dpex_exp_kernel_target	1✔
306	_fold_args = False	1✔
307
308	def __init__(	1✔
309	self,
310	pyfunc,
311	local_vars_to_numba_types=None,
312	targetoptions=None,
313	pipeline_class=kernel_compiler.KernelCompiler,
314	):
315	if targetoptions is None:	1!
316	targetoptions = {}	×
317
318	if local_vars_to_numba_types is None:	1!
319	local_vars_to_numba_types = {}	1✔
320
321	targetoptions["nopython"] = True	1✔
322	targetoptions["experimental"] = True	1✔
323
324	self._kernel_name = pyfunc.__name__	1✔
325
326	if numba_sem_version < (0, 59, 0):	1!
327	# pylint: disable=unexpected-keyword-arg
328	super().__init__(	1✔
329	py_func=pyfunc,
330	locals=local_vars_to_numba_types,
331	impl_kind="direct",
332	targetoptions=targetoptions,
333	pipeline_class=pipeline_class,
334	)
335	else:
336	super().__init__(	×
337	py_func=pyfunc,
338	locals=local_vars_to_numba_types,
339	targetoptions=targetoptions,
340	pipeline_class=pipeline_class,
341	)
342	self._compiler = _SPIRVKernelCompiler(	1✔
343	pyfunc,
344	self.targetdescr,
345	targetoptions,
346	local_vars_to_numba_types,
347	pipeline_class,
348	)
349
350	def typeof_pyval(self, val):	1✔
351	"""
352	Resolve the Numba type of Python value val.
353	This is called from numba._dispatcher as a fallback if the native code
354	cannot decide the type.
355	"""
356	# Not going through the resolve_argument_type() indirection
357	# can save a couple µs.
358	try:	×
359	typ = typeof(val, Purpose.argument)	×
360	if isinstance(typ, types.Array) and not isinstance(typ, USMNdArray):	×
361	raise UnsupportedKernelArgumentError(	×
362	type=str(type(val)), value=val
363	)
364	except ValueError:	×
365	typ = types.pyobject	×
366	else:
367	if typ is None:	×
368	typ = types.pyobject	×
369	self._types_active_call.append(typ)	×
370	return typ	×
371
372	def add_overload(self, cres):	1✔
373	args = tuple(cres.signature.args)	1✔
374	self.overloads[args] = cres	1✔
375
376	def compile(self, sig) -> any:	1✔
377	disp = self._get_dispatcher_for_current_target()	1✔
378	if disp is not self:	1!
379	return disp.compile(sig)	×
380
381	with ExitStack() as scope:	1✔
382	cres = None	1✔
383
384	def cb_compiler(dur):	1✔
385	if cres is not None:	1!
386	self._callback_add_compiler_timer(dur, cres)	×
387
388	def cb_llvm(dur):	1✔
389	if cres is not None:	1!
390	self._callback_add_llvm_timer(dur, cres)	×
391
392	scope.enter_context(	1✔
393	ev.install_timer("numba:compiler_lock", cb_compiler)
394	)
395	scope.enter_context(ev.install_timer("numba:llvm_lock", cb_llvm))	1✔
396	scope.enter_context(global_compiler_lock)	1✔
397
398	if not self._can_compile:	1!
399	raise RuntimeError("compilation disabled")	×
400	# Use counter to track recursion compilation depth
401	with self._compiling_counter:	1✔
402	args, return_type = sigutils.normalize_signature(sig)	1✔
403
404	if (	1✔
405	self.targetoptions["_compilation_mode"]
406	== CompilationMode.KERNEL
407	):
408	# Compute follows data based queue equivalence is only
409	# evaluated for kernel functions whose arguments are
410	# supposed to be arrays. For device_func decorated
411	# functions, the arguments can be scalar and we skip queue
412	# equivalence check.
413	try:	1✔
414	self._compiler.check_arguments(self._kernel_name, args)	1✔
415	except ExecutionQueueInferenceError as eqie:	1✔
416	raise eqie	1✔
417
418	# A function being compiled in the KERNEL compilation mode
419	# cannot have a non-void return value
420	if return_type and return_type != void:	1!
421	raise KernelHasReturnValueError(	×
422	kernel_name=None, return_type=return_type, sig=sig
423	)
424
425	# Don't recompile if signature already exists
426	existing = self.overloads.get(tuple(args))	1✔
427	if existing is not None:	1✔
428	return existing.entry_point	1✔
429
430	# TODO: Enable caching
431	# Add code to enable on disk caching of a binary spirv kernel.
432	# Refer: https://github.com/IntelPython/numba-dpex/issues/1197
433	self._cache_misses[sig] += 1	1✔
434	with ev.trigger_event(	1✔
435	"numba_dpex:compile",
436	data={
437	"dispatcher": self,
438	"args": args,
439	"return_type": return_type,
440	},
441	):
442	try:	1✔
443	compiler: _SPIRVKernelCompiler = self._compiler	1✔
444	kcres: _SPIRVKernelCompileResult = compiler.compile(	1✔
445	args, return_type
446	)
447	except errors.ForceLiteralArg as err:	1✔
448
449	def folded(args, kws):	×
450	return self._compiler.fold_argument_types(	×
451	args, kws
452	)[1]
453
454	raise err.bind_fold_arguments(folded)	×
455	self.add_overload(kcres)	1✔
456
457	kcres.target_context.insert_user_function(	1✔
458	kcres.entry_point, kcres.fndesc, [kcres.library]
459	)
460
461	# TODO: enable caching of kernel_module
462	# https://github.com/IntelPython/numba-dpex/issues/1197
463
464	return kcres.entry_point	1✔
465
466	def __getitem__(self, args):	1✔
467	"""Square-bracket notation for configuring launch arguments is not
468	supported.
469	"""
470
471	raise NotImplementedError
472
473	def __call__(self, args, *kw_args):	1✔
474	"""Functor to launch a kernel."""
475
476	raise NotImplementedError
477
478
479	_dpex_target = target_registry[DPEX_KERNEL_EXP_TARGET_NAME]	1✔
480	dispatcher_registry[_dpex_target] = SPIRVKernelDispatcher	1✔

IntelPython / numba-dpex / 8289101957

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