8350440097

Committed 19 Mar 2024 09:37PM UTC coverage: 83.04% (+0.3%) from 82.76%

Build # 8350440097

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1331 from IntelPython/experimental/local_accessors

A sycl::local_accessor-like API for numba-dpex kernel

Run Details

1626 of 2227 branches covered (73.01%)

Branch coverage included in aggregate %.

138 of 153 new or added lines in 12 files covered. (90.2%)

1 existing line in 1 file now uncovered.

6634 of 7720 relevant lines covered (85.93%)

0.86 hits per line

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

96.69

/numba_dpex/experimental/launcher.py

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

"""Provides a helper function to call a numba_dpex.kernel decorated function
from either CPython or a numba_dpex.dpjit decorated function.
"""

import warnings
from inspect import signature
from typing import NamedTuple, Union

import dpctl
from llvmlite import ir as llvmir
from numba.core import cgutils, types
from numba.core.cpu import CPUContext
from numba.core.types.containers import Tuple, UniTuple
from numba.core.types.functions import Dispatcher
from numba.extending import intrinsic

from numba_dpex import dpjit
from numba_dpex.core.targets.dpjit_target import DPEX_TARGET_NAME
from numba_dpex.core.types import DpctlSyclEvent, NdRangeType, RangeType
from numba_dpex.core.types.kernel_api.index_space_ids import (
    ItemType,
    NdItemType,
)
from numba_dpex.core.types.kernel_api.local_accessor import LocalAccessorType
from numba_dpex.core.utils import kernel_launcher as kl
from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
from numba_dpex.dpctl_iface.wrappers import wrap_event_reference
from numba_dpex.kernel_api_impl.spirv.dispatcher import (
    SPIRVKernelDispatcher,
    _SPIRVKernelCompileResult,
)
from numba_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext


class _LLRange(NamedTuple):
    """Analog of Range and NdRange but for the llvm ir values."""

    global_range_extents: list
    local_range_extents: list


def _has_a_local_accessor_argument(args):
    """Checks if there exists at least one LocalAccessorType object in the
    input tuple.

    Args:
        args (_type_): A tuple of numba.core.Type objects

    Returns:
        bool : True if at least one LocalAccessorType object was found,
            otherwise False.
    """
    for arg in args:
        if isinstance(arg, LocalAccessorType):
            return True
    return False


def _wrap_event_reference_tuple(ctx, builder, event1, event2):
    """Creates tuple data model from two event data models, so it can be
    boxed to Python."""
    ty_event = DpctlSyclEvent()
    tupty = types.Tuple([ty_event, ty_event])

    lltupty = ctx.get_value_type(tupty)
    tup = cgutils.get_null_value(lltupty)
    tup = builder.insert_value(tup, event1, 0)
    tup = builder.insert_value(tup, event2, 1)

    return tup


@intrinsic(target=DPEX_TARGET_NAME)
def _submit_kernel_async(
    typingctx,
    ty_kernel_fn: Dispatcher,
    ty_index_space: Union[RangeType, NdRangeType],
    ty_dependent_events: UniTuple,
    ty_kernel_args_tuple: UniTuple,
):
    """Generates IR code for call_kernel_async ``dpjit`` function."""
    return _submit_kernel(
        typingctx,
        ty_kernel_fn,
        ty_index_space,
        ty_dependent_events,
        ty_kernel_args_tuple,
        sync=False,
    )


@intrinsic(target=DPEX_TARGET_NAME)
def _submit_kernel_sync(
    typingctx,
    ty_kernel_fn: Dispatcher,
    ty_index_space: Union[RangeType, NdRangeType],
    ty_kernel_args_tuple: UniTuple,
):
    """Generates IR code for call_kernel ``dpjit`` function."""
    return _submit_kernel(
        typingctx,
        ty_kernel_fn,
        ty_index_space,
        None,
        ty_kernel_args_tuple,
        sync=True,
    )


def _submit_kernel(  # pylint: disable=too-many-arguments
    typingctx,  # pylint: disable=unused-argument
    ty_kernel_fn: Dispatcher,
    ty_index_space: Union[RangeType, NdRangeType],
    ty_dependent_events: UniTuple,
    ty_kernel_args_tuple: UniTuple,
    sync: bool,
):
    """Generates IR code for call_kernel_{async|sync} ``dpjit`` function.

    The intrinsic first compiles the kernel function to SPIR-V, and then to a
    SYCL kernel bundle. The arguments to the kernel are also packed into
    flattened arrays and the SYCL queue to which the kernel will be submitted
    extracted from the args. Finally, the actual kernel is extracted from the
    kernel bundle and submitted to the SYCL queue.

    If sync set to False, it acquires memory infos from kernel arguments to
    prevent garbage collection on them. Then it schedules host task to release
    that arguments and unblock garbage collection. Tuple of host task and device
    tasks are returned.
    """
    # signature of this intrinsic
    ty_return = types.void
    if not sync:
        ty_event = DpctlSyclEvent()
        ty_return = types.Tuple([ty_event, ty_event])

    if ty_dependent_events is not None:
        if not isinstance(ty_dependent_events, UniTuple) and not isinstance(
            ty_dependent_events, Tuple
        ):
            raise ValueError("dependent events must be passed as a tuple")

        sig = ty_return(
            ty_kernel_fn,
            ty_index_space,
            ty_dependent_events,
            ty_kernel_args_tuple,
        )
    else:
        sig = ty_return(ty_kernel_fn, ty_index_space, ty_kernel_args_tuple)

    # Add Item/NdItem as a first argument to kernel arguments list. It is
    # an empty struct so any other modifications at kernel submission are not
    # needed.
    if len(signature(ty_kernel_fn.dispatcher.py_func).parameters) > len(
        ty_kernel_args_tuple
    ):
        if isinstance(ty_index_space, RangeType):
            ty_item = ItemType(ty_index_space.ndim)
        else:
            ty_item = NdItemType(ty_index_space.ndim)

        ty_kernel_args_tuple = (ty_item, *ty_kernel_args_tuple)
    else:
        warnings.warn(
            "Kernels without item/nd_item will be not supported in the future",
            DeprecationWarning,
        )

    # Validate local accessor arguments are passed only to a kernel that is
    # launched with an NdRange index space. Reference section 4.7.6.11. of the
    # SYCL 2020 specification: A local_accessor must not be used in a SYCL
    # kernel function that is invoked via single_task or via the simple form of
    # parallel_for that takes a range parameter.
    if _has_a_local_accessor_argument(ty_kernel_args_tuple) and isinstance(
        ty_index_space, RangeType
    ):
        raise TypeError(
            "A RangeType kernel cannot have a LocalAccessor argument"
        )

    # ty_kernel_fn is type specific to exact function, so we can get function
    # directly from type and compile it. Thats why we don't need to get it in
    # codegen
    kernel_dispatcher: SPIRVKernelDispatcher = ty_kernel_fn.dispatcher
    kcres: _SPIRVKernelCompileResult = kernel_dispatcher.get_compile_result(
        types.void(*ty_kernel_args_tuple)  # kernel signature
    )
    kernel_module: kl.SPIRVKernelModule = kcres.kernel_device_ir_module
    kernel_targetctx: SPIRVTargetContext = kernel_dispatcher.targetctx

    def codegen(
        cgctx: CPUContext, builder: llvmir.IRBuilder, sig, llargs: list
    ):
        ty_index_space: Union[RangeType, NdRangeType] = sig.args[1]
        ll_index_space: llvmir.Instruction = llargs[1]
        ty_kernel_args_tuple: UniTuple = sig.args[-1]
        ll_kernel_args_tuple: llvmir.Instruction = llargs[-1]

        if len(llargs) == 4:
            ty_dependent_events: UniTuple = sig.args[2]
            ll_dependent_events: llvmir.Instruction = llargs[2]
        else:
            ty_dependent_events = None

        kl_builder = kl.KernelLaunchIRBuilder(
            cgctx,
            builder,
            kernel_targetctx.data_model_manager,
        )
        kl_builder.set_range_from_indexer(
            ty_indexer_arg=ty_index_space,
            ll_index_arg=ll_index_space,
        )
        kl_builder.set_arguments_form_tuple(
            ty_kernel_args_tuple, ll_kernel_args_tuple
        )
        kl_builder.set_queue_from_arguments()
        kl_builder.set_kernel_from_spirv(kernel_module)
        if ty_dependent_events is None:
            kl_builder.set_dependent_events([])
        else:
            kl_builder.set_dependent_events_from_tuple(
                ty_dependent_events,
                ll_dependent_events,
            )
        device_event_ref = kl_builder.submit()

        if not sync:
            host_event_ref = kl_builder.acquire_meminfo_and_submit_release()

            return _wrap_event_reference_tuple(
                cgctx,
                builder,
                wrap_event_reference(cgctx, builder, host_event_ref),
                wrap_event_reference(cgctx, builder, device_event_ref),
            )

        sycl.dpctl_event_wait(builder, device_event_ref)
        sycl.dpctl_event_delete(builder, device_event_ref)

        return None

    return sig, codegen


@dpjit
def call_kernel(kernel_fn, index_space, *kernel_args) -> None:
    """Compiles and synchronously executes a kernel function.

    Kernel execution happens in synchronous way, so the main thread will be
    blocked till the kernel done execution.

    Args:
        kernel_fn (numba_dpex.experimental.KernelDispatcher): A
            :func:`numba_dpex.experimental.kernel` decorated function that is
            compiled to a ``KernelDispatcher``.
        index_space (Range | NdRange): A Range or NdRange type object that
            specifies the index space for the kernel.
        kernel_args : List of objects that are passed to the numba_dpex.kernel
            decorated function.
    """
    _submit_kernel_sync(  # pylint: disable=E1120
        kernel_fn,
        index_space,
        kernel_args,
    )


@dpjit
def call_kernel_async(
    kernel_fn,
    index_space,
    dependent_events: list[dpctl.SyclEvent],
    *kernel_args
) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]:
    """Compiles and asynchronously executes a kernel function.

    Calls a :func:`numba_dpex.experimental.kernel` decorated function
    asynchronously from CPython or from a :func:`numba_dpex.dpjit` function. As
    the kernel execution happens asynchronously, so the main thread will not be
    blocked till the kernel done execution. Instead the function returns back to
    caller a handle for an *event* to track kernel execution. It is a user's
    responsibility to properly track kernel execution completion and not use any
    data that may still be used by the kernel prior to the kernel's completion.

    Args:
        kernel_fn (KernelDispatcher): A
            :func:`numba_dpex.experimental.kernel` decorated function that is
            compiled to a ``KernelDispatcher``.
        index_space (Range | NdRange): A Range or NdRange type object that
            specifies the index space for the kernel.
        kernel_args : List of objects that are passed to the numba_dpex.kernel
            decorated function.

    Returns:
        A pair of ``dpctl.SyclEvent`` objects. The pair of events constitute of
        a host task and an event associated with the kernel execution. The event
        associated with the kernel execution indicates the execution status of
        the submitted kernel function. The host task manages the lifetime of any
        PyObject passed in as a kernel argument and automatically decrements the
        reference count of the object on kernel execution completion.
    """
    return _submit_kernel_async(  # pylint: disable=E1120
        kernel_fn,
        index_space,
        dependent_events,
        kernel_args,
    )

1	# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	"""Provides a helper function to call a numba_dpex.kernel decorated function	1✔
6	from either CPython or a numba_dpex.dpjit decorated function.
7	"""
8
9	import warnings	1✔
10	from inspect import signature	1✔
11	from typing import NamedTuple, Union	1✔
12
13	import dpctl	1✔
14	from llvmlite import ir as llvmir	1✔
15	from numba.core import cgutils, types	1✔
16	from numba.core.cpu import CPUContext	1✔
17	from numba.core.types.containers import Tuple, UniTuple	1✔
18	from numba.core.types.functions import Dispatcher	1✔
19	from numba.extending import intrinsic	1✔
20
21	from numba_dpex import dpjit	1✔
22	from numba_dpex.core.targets.dpjit_target import DPEX_TARGET_NAME	1✔
23	from numba_dpex.core.types import DpctlSyclEvent, NdRangeType, RangeType	1✔
24	from numba_dpex.core.types.kernel_api.index_space_ids import (	1✔
25	ItemType,
26	NdItemType,
27	)
28	from numba_dpex.core.types.kernel_api.local_accessor import LocalAccessorType	1✔
29	from numba_dpex.core.utils import kernel_launcher as kl	1✔
30	from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl	1✔
31	from numba_dpex.dpctl_iface.wrappers import wrap_event_reference	1✔
32	from numba_dpex.kernel_api_impl.spirv.dispatcher import (	1✔
33	SPIRVKernelDispatcher,
34	_SPIRVKernelCompileResult,
35	)
36	from numba_dpex.kernel_api_impl.spirv.target import SPIRVTargetContext	1✔
37
38
39	class _LLRange(NamedTuple):	1✔
40	"""Analog of Range and NdRange but for the llvm ir values."""
41
42	global_range_extents: list	1✔
43	local_range_extents: list	1✔
44
45
46	def _has_a_local_accessor_argument(args):	1✔
47	"""Checks if there exists at least one LocalAccessorType object in the
48	input tuple.
49
50	Args:
51	args (_type_): A tuple of numba.core.Type objects
52
53	Returns:
54	bool : True if at least one LocalAccessorType object was found,
55	otherwise False.
56	"""
57	for arg in args:	1✔
58	if isinstance(arg, LocalAccessorType):	1✔
59	return True	1✔
60	return False	1✔
61
62
63	def _wrap_event_reference_tuple(ctx, builder, event1, event2):	1✔
64	"""Creates tuple data model from two event data models, so it can be
65	boxed to Python."""
66	ty_event = DpctlSyclEvent()	1✔
67	tupty = types.Tuple([ty_event, ty_event])	1✔
68
69	lltupty = ctx.get_value_type(tupty)	1✔
70	tup = cgutils.get_null_value(lltupty)	1✔
71	tup = builder.insert_value(tup, event1, 0)	1✔
72	tup = builder.insert_value(tup, event2, 1)	1✔
73
74	return tup	1✔
75
76
77	@intrinsic(target=DPEX_TARGET_NAME)	1✔
78	def _submit_kernel_async(	1✔
79	typingctx,
80	ty_kernel_fn: Dispatcher,
81	ty_index_space: Union[RangeType, NdRangeType],
82	ty_dependent_events: UniTuple,
83	ty_kernel_args_tuple: UniTuple,
84	):
85	"""Generates IR code for call_kernel_async ``dpjit`` function."""
86	return _submit_kernel(	1✔
87	typingctx,
88	ty_kernel_fn,
89	ty_index_space,
90	ty_dependent_events,
91	ty_kernel_args_tuple,
92	sync=False,
93	)
94
95
96	@intrinsic(target=DPEX_TARGET_NAME)	1✔
97	def _submit_kernel_sync(	1✔
98	typingctx,
99	ty_kernel_fn: Dispatcher,
100	ty_index_space: Union[RangeType, NdRangeType],
101	ty_kernel_args_tuple: UniTuple,
102	):
103	"""Generates IR code for call_kernel ``dpjit`` function."""
104	return _submit_kernel(	1✔
105	typingctx,
106	ty_kernel_fn,
107	ty_index_space,
108	None,
109	ty_kernel_args_tuple,
110	sync=True,
111	)
112
113
114	def _submit_kernel( # pylint: disable=too-many-arguments	1✔
115	typingctx, # pylint: disable=unused-argument
116	ty_kernel_fn: Dispatcher,
117	ty_index_space: Union[RangeType, NdRangeType],
118	ty_dependent_events: UniTuple,
119	ty_kernel_args_tuple: UniTuple,
120	sync: bool,
121	):
122	"""Generates IR code for call_kernel_{async\|sync} ``dpjit`` function.
123
124	The intrinsic first compiles the kernel function to SPIR-V, and then to a
125	SYCL kernel bundle. The arguments to the kernel are also packed into
126	flattened arrays and the SYCL queue to which the kernel will be submitted
127	extracted from the args. Finally, the actual kernel is extracted from the
128	kernel bundle and submitted to the SYCL queue.
129
130	If sync set to False, it acquires memory infos from kernel arguments to
131	prevent garbage collection on them. Then it schedules host task to release
132	that arguments and unblock garbage collection. Tuple of host task and device
133	tasks are returned.
134	"""
135	# signature of this intrinsic
136	ty_return = types.void	1✔
137	if not sync:	1✔
138	ty_event = DpctlSyclEvent()	1✔
139	ty_return = types.Tuple([ty_event, ty_event])	1✔
140
141	if ty_dependent_events is not None:	1✔
142	if not isinstance(ty_dependent_events, UniTuple) and not isinstance(	1✔
143	ty_dependent_events, Tuple
144	):
145	raise ValueError("dependent events must be passed as a tuple")	1✔
146
147	sig = ty_return(	1✔
148	ty_kernel_fn,
149	ty_index_space,
150	ty_dependent_events,
151	ty_kernel_args_tuple,
152	)
153	else:
154	sig = ty_return(ty_kernel_fn, ty_index_space, ty_kernel_args_tuple)	1✔
155
156	# Add Item/NdItem as a first argument to kernel arguments list. It is
157	# an empty struct so any other modifications at kernel submission are not
158	# needed.
159	if len(signature(ty_kernel_fn.dispatcher.py_func).parameters) > len(	1✔
160	ty_kernel_args_tuple
161	):
162	if isinstance(ty_index_space, RangeType):	1✔
163	ty_item = ItemType(ty_index_space.ndim)	1✔
164	else:
165	ty_item = NdItemType(ty_index_space.ndim)	1✔
166
167	ty_kernel_args_tuple = (ty_item, *ty_kernel_args_tuple)	1✔
168	else:
169	warnings.warn(	1✔
170	"Kernels without item/nd_item will be not supported in the future",
171	DeprecationWarning,
172	)
173
174	# Validate local accessor arguments are passed only to a kernel that is
175	# launched with an NdRange index space. Reference section 4.7.6.11. of the
176	# SYCL 2020 specification: A local_accessor must not be used in a SYCL
177	# kernel function that is invoked via single_task or via the simple form of
178	# parallel_for that takes a range parameter.
179	if _has_a_local_accessor_argument(ty_kernel_args_tuple) and isinstance(	1!
180	ty_index_space, RangeType
181	):
NEW 182	raise TypeError(	×
183	"A RangeType kernel cannot have a LocalAccessor argument"
184	)
185
186	# ty_kernel_fn is type specific to exact function, so we can get function
187	# directly from type and compile it. Thats why we don't need to get it in
188	# codegen
189	kernel_dispatcher: SPIRVKernelDispatcher = ty_kernel_fn.dispatcher	1✔
190	kcres: _SPIRVKernelCompileResult = kernel_dispatcher.get_compile_result(	1✔
191	types.void(*ty_kernel_args_tuple) # kernel signature
192	)
193	kernel_module: kl.SPIRVKernelModule = kcres.kernel_device_ir_module	1✔
194	kernel_targetctx: SPIRVTargetContext = kernel_dispatcher.targetctx	1✔
195
196	def codegen(	1✔
197	cgctx: CPUContext, builder: llvmir.IRBuilder, sig, llargs: list
198	):
199	ty_index_space: Union[RangeType, NdRangeType] = sig.args[1]	1✔
200	ll_index_space: llvmir.Instruction = llargs[1]	1✔
201	ty_kernel_args_tuple: UniTuple = sig.args[-1]	1✔
202	ll_kernel_args_tuple: llvmir.Instruction = llargs[-1]	1✔
203
204	if len(llargs) == 4:	1✔
205	ty_dependent_events: UniTuple = sig.args[2]	1✔
206	ll_dependent_events: llvmir.Instruction = llargs[2]	1✔
207	else:
208	ty_dependent_events = None	1✔
209
210	kl_builder = kl.KernelLaunchIRBuilder(	1✔
211	cgctx,
212	builder,
213	kernel_targetctx.data_model_manager,
214	)
215	kl_builder.set_range_from_indexer(	1✔
216	ty_indexer_arg=ty_index_space,
217	ll_index_arg=ll_index_space,
218	)
219	kl_builder.set_arguments_form_tuple(	1✔
220	ty_kernel_args_tuple, ll_kernel_args_tuple
221	)
222	kl_builder.set_queue_from_arguments()	1✔
223	kl_builder.set_kernel_from_spirv(kernel_module)	1✔
224	if ty_dependent_events is None:	1✔
225	kl_builder.set_dependent_events([])	1✔
226	else:
227	kl_builder.set_dependent_events_from_tuple(	1✔
228	ty_dependent_events,
229	ll_dependent_events,
230	)
231	device_event_ref = kl_builder.submit()	1✔
232
233	if not sync:	1✔
234	host_event_ref = kl_builder.acquire_meminfo_and_submit_release()	1✔
235
236	return _wrap_event_reference_tuple(	1✔
237	cgctx,
238	builder,
239	wrap_event_reference(cgctx, builder, host_event_ref),
240	wrap_event_reference(cgctx, builder, device_event_ref),
241	)
242
243	sycl.dpctl_event_wait(builder, device_event_ref)	1✔
244	sycl.dpctl_event_delete(builder, device_event_ref)	1✔
245
246	return None	1✔
247
248	return sig, codegen	1✔
249
250
251	@dpjit	1✔
252	def call_kernel(kernel_fn, index_space, *kernel_args) -> None:	1✔
253	"""Compiles and synchronously executes a kernel function.
254
255	Kernel execution happens in synchronous way, so the main thread will be
256	blocked till the kernel done execution.
257
258	Args:
259	kernel_fn (numba_dpex.experimental.KernelDispatcher): A
260	:func:`numba_dpex.experimental.kernel` decorated function that is
261	compiled to a ``KernelDispatcher``.
262	index_space (Range \| NdRange): A Range or NdRange type object that
263	specifies the index space for the kernel.
264	kernel_args : List of objects that are passed to the numba_dpex.kernel
265	decorated function.
266	"""
267	_submit_kernel_sync( # pylint: disable=E1120	×
268	kernel_fn,
269	index_space,
270	kernel_args,
271	)
272
273
274	@dpjit	1✔
275	def call_kernel_async(	1✔
276	kernel_fn,
277	index_space,
278	dependent_events: list[dpctl.SyclEvent],
279	*kernel_args
280	) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]:
281	"""Compiles and asynchronously executes a kernel function.
282
283	Calls a :func:`numba_dpex.experimental.kernel` decorated function
284	asynchronously from CPython or from a :func:`numba_dpex.dpjit` function. As
285	the kernel execution happens asynchronously, so the main thread will not be
286	blocked till the kernel done execution. Instead the function returns back to
287	caller a handle for an event to track kernel execution. It is a user's
288	responsibility to properly track kernel execution completion and not use any
289	data that may still be used by the kernel prior to the kernel's completion.
290
291	Args:
292	kernel_fn (KernelDispatcher): A
293	:func:`numba_dpex.experimental.kernel` decorated function that is
294	compiled to a ``KernelDispatcher``.
295	index_space (Range \| NdRange): A Range or NdRange type object that
296	specifies the index space for the kernel.
297	kernel_args : List of objects that are passed to the numba_dpex.kernel
298	decorated function.
299
300	Returns:
301	A pair of ``dpctl.SyclEvent`` objects. The pair of events constitute of
302	a host task and an event associated with the kernel execution. The event
303	associated with the kernel execution indicates the execution status of
304	the submitted kernel function. The host task manages the lifetime of any
305	PyObject passed in as a kernel argument and automatically decrements the
306	reference count of the object on kernel execution completion.
307	"""
308	return _submit_kernel_async( # pylint: disable=E1120	×
309	kernel_fn,
310	index_space,
311	dependent_events,
312	kernel_args,
313	)

IntelPython / numba-dpex / 8350440097

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