7199556298

Committed 13 Dec 2023 06:32PM UTC coverage: 79.309% (+0.07%) from 79.239%

Build # 7199556298

Build Type

Pull #1243

github

Committed by

web-flow

Commit Message

Merge 12dbc5ae4 into 8fc7fbddb

Pull Request Pull Request #1243: Apply builder pattern to kernel launcher

Run Details

1197 of 1769 branches covered (0.0%)

Branch coverage included in aggregate %.

104 of 105 new or added lines in 5 files covered. (99.05%)

25 existing lines in 4 files now uncovered.

5323 of 6452 relevant lines covered (82.5%)

0.83 hits per line

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

92.76

/numba_dpex/experimental/launcher.py

# SPDX-FileCopyrightText: 2023 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.
"""

from typing import NamedTuple, Union

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

from numba_dpex import config, dpjit, utils
from numba_dpex.core.exceptions import UnreachableError
from numba_dpex.core.runtime.context import DpexRTContext
from numba_dpex.core.types import (
    DpctlSyclEvent,
    DpnpNdArray,
    NdRangeType,
    RangeType,
)
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.experimental.kernel_dispatcher import (
    KernelDispatcher,
    _KernelModule,
)
from numba_dpex.utils import create_null_ptr


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

    global_range_extents: list
    local_range_extents: list


class CallKernelBuilder:
    """
    Helper class to build LLVM IR for a numba function that calls kernel.
    """

    def __init__(
        self,
        codegen_targetctx: cpu.CPUContext,
        builder: llvmir.IRBuilder,
    ):
        self.context = codegen_targetctx
        self.builder = builder
        # TODO: get dpex RT from cached property once the PR is merged
        # https://github.com/IntelPython/numba-dpex/pull/1027
        # and get rid of the global variable. Use self.context.dpexrt instead.
        self.dpexrt = DpexRTContext(self.context)

        if config.DEBUG_KERNEL_LAUNCHER:
            cgutils.printf(
                self.builder,
                "DPEX-DEBUG: Inside the kernel launcher function\n",
            )

    def extract_arguments_from_tuple(
        self,
        ty_kernel_args_tuple: UniTuple,
        ll_kernel_args_tuple: llvmir.Instruction,
    ) -> list[llvmir.Instruction]:
        """Extracts LLVM IR values from llvm tuple into python array."""

        kernel_args = []
        for pos in range(len(ty_kernel_args_tuple)):
            kernel_args.append(
                self.builder.extract_value(ll_kernel_args_tuple, pos)
            )

        return kernel_args

    def allocate_meminfo_array(
        self,
        kernel_argtys: tuple[types.Type, ...],
        kernel_args: [llvmir.Instruction, ...],
    ) -> tuple[int, list[llvmir.Instruction]]:
        """Allocates an LLVM array value to store each memory info from all
        kernel arguments. The array is the populated with the LLVM value for
        every meminfo of the kernel arguments.
        """
        meminfos = []
        for arg_num, argtype in enumerate(kernel_argtys):
            llvm_val = kernel_args[arg_num]

            meminfos += [
                meminfo
                for ty, meminfo in self.context.nrt.get_meminfos(
                    self.builder, argtype, llvm_val
                )
            ]

        meminfo_list = cgutils.alloca_once(
            self.builder,
            utils.get_llvm_type(context=self.context, type=types.voidptr),
            size=self.context.get_constant(types.uintp, len(meminfos)),
        )

        for meminfo_num, meminfo in enumerate(meminfos):
            meminfo_arg_dst = self.builder.gep(
                meminfo_list,
                [self.context.get_constant(types.int32, meminfo_num)],
            )
            meminfo_ptr = self.builder.bitcast(
                meminfo,
                utils.get_llvm_type(context=self.context, type=types.voidptr),
            )
            self.builder.store(meminfo_ptr, meminfo_arg_dst)

        return len(meminfos), meminfo_list

    def get_queue_ref_val(
        self,
        kernel_argtys: tuple[types.Type, ...],
        kernel_args: [llvmir.Instruction, ...],
    ):
        """
        Get the sycl queue from the first DpnpNdArray argument. Prior passes
        before lowering make sure that compute-follows-data is enforceable
        for a specific call to a kernel. As such, at the stage of lowering
        the queue from the first DpnpNdArray argument can be extracted.
        """

        for arg_num, argty in enumerate(kernel_argtys):
            if isinstance(argty, DpnpNdArray):
                llvm_val = kernel_args[arg_num]
                datamodel = self.context.data_model_manager.lookup(argty)
                sycl_queue_attr_pos = datamodel.get_field_position("sycl_queue")
                ptr_to_queue_ref = self.builder.extract_value(
                    llvm_val, sycl_queue_attr_pos
                )
                break

        return ptr_to_queue_ref

    def get_kernel(self, queue_ref, kernel_module: _KernelModule):
        """Returns the pointer to the sycl::kernel object in a passed in
        sycl::kernel_bundle wrapper object.
        """
        # Inserts a global constant byte string in the current LLVM module to
        # store the passed in SPIR-V binary blob.
        kernel_bc_byte_str = self.context.insert_const_bytes(
            self.builder.module,
            bytes=kernel_module.kernel_bitcode,
        )

        kernel_name = self.context.insert_const_string(
            self.builder.module, kernel_module.kernel_name
        )

        context_ref = sycl.dpctl_queue_get_context(self.builder, queue_ref)
        device_ref = sycl.dpctl_queue_get_device(self.builder, queue_ref)

        # build_or_get_kernel stills reference to context and device cause it
        # needs to keep them alive for keys.
        kernel_ref = self.dpexrt.build_or_get_kernel(
            self.builder,
            [
                context_ref,
                device_ref,
                llvmir.Constant(
                    llvmir.IntType(64), hash(kernel_module.kernel_bitcode)
                ),
                kernel_bc_byte_str,
                llvmir.Constant(
                    llvmir.IntType(64), len(kernel_module.kernel_bitcode)
                ),
                self.builder.load(create_null_ptr(self.builder, self.context)),
                kernel_name,
            ],
        )

        return kernel_ref

    def create_llvm_values_for_index_space(
        self,
        ty_indexer_arg: Union[RangeType, NdRangeType],
        index_arg: llvmir.BaseStructType,
    ) -> LLRange:
        """Returns two lists of LLVM IR Values that hold the unboxed extents of
        a Python Range or NdRange object.
        """
        ndim = ty_indexer_arg.ndim
        global_range_extents = []
        local_range_extents = []
        indexer_datamodel = numba_default_dmm.lookup(ty_indexer_arg)

        if isinstance(ty_indexer_arg, RangeType):
            for dim_num in range(ndim):
                dim_pos = indexer_datamodel.get_field_position(
                    "dim" + str(dim_num)
                )
                global_range_extents.append(
                    self.builder.extract_value(index_arg, dim_pos)
                )
        elif isinstance(ty_indexer_arg, NdRangeType):
            for dim_num in range(ndim):
                gdim_pos = indexer_datamodel.get_field_position(
                    "gdim" + str(dim_num)
                )
                global_range_extents.append(
                    self.builder.extract_value(index_arg, gdim_pos)
                )
                ldim_pos = indexer_datamodel.get_field_position(
                    "ldim" + str(dim_num)
                )
                local_range_extents.append(
                    self.builder.extract_value(index_arg, ldim_pos)
                )
        else:
            raise UnreachableError

        return LLRange(global_range_extents, local_range_extents)

    def acquire_meminfo_and_schedule_release(
        self,
        queue_ref,
        event_ref,
        ty_kernel_args,
        kernel_args,
    ):
        """Schedule sycl host task to release nrt meminfo of the arguments used
        to run job. Use it to keep arguments alive during kernel execution."""
        total_meminfos, meminfo_list = self.allocate_meminfo_array(
            ty_kernel_args, kernel_args
        )

        event_ref_ptr = self.builder.alloca(event_ref.type)
        self.builder.store(event_ref, event_ref_ptr)

        status_ptr = cgutils.alloca_once(
            self.builder, self.context.get_value_type(types.uint64)
        )
        host_eref = self.dpexrt.acquire_meminfo_and_schedule_release(
            self.builder,
            [
                self.context.nrt.get_nrt_api(self.builder),
                queue_ref,
                meminfo_list,
                self.context.get_constant(types.uintp, total_meminfos),
                event_ref_ptr,
                self.context.get_constant(types.uintp, 1),
                status_ptr,
            ],
        )
        return host_eref


@intrinsic(target="cpu")
def _submit_kernel(
    typingctx,  # pylint: disable=W0613
    ty_kernel_fn: Dispatcher,
    ty_index_space: Union[RangeType, NdRangeType],
    ty_kernel_args_tuple: UniTuple,
):
    """Generates IR code for call_kernel dpjit function.

    The intrinsic first compiles the kernel function to SPIRV, 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.
    """
    # signature of this intrinsic
    ty_event = DpctlSyclEvent()
    sig = ty_event(ty_kernel_fn, ty_index_space, ty_kernel_args_tuple)
    kernel_sig = types.void(*ty_kernel_args_tuple)
    # 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: KernelDispatcher = ty_kernel_fn.dispatcher
    kernel_dispatcher.compile(kernel_sig)
    kernel_module: _KernelModule = kernel_dispatcher.get_overload_kcres(
        kernel_sig
    ).kernel_device_ir_module
    kernel_targetctx = kernel_dispatcher.targetctx

    def codegen(cgctx, builder, sig, llargs):
        # llargs[0] is kernel function that we don't need anymore (see above)
        ty_index_space: Union[RangeType, NdRangeType] = sig.args[1]
        ll_index_space: llvmir.Instruction = llargs[1]
        ty_kernel_args_tuple: UniTuple = sig.args[2]
        ll_kernel_args_tuple: llvmir.Instruction = llargs[2]

        generator = CallKernelBuilder(
            codegen_targetctx=cgctx,
            builder=builder,
        )

        kernel_args = generator.extract_arguments_from_tuple(
            ty_kernel_args_tuple=ty_kernel_args_tuple,
            ll_kernel_args_tuple=ll_kernel_args_tuple,
        )

        # queue_ref is just a pointer to the attribute, so we don't have to
        # clean it up
        queue_ref = generator.get_queue_ref_val(
            kernel_argtys=ty_kernel_args_tuple,
            kernel_args=kernel_args,
        )

        # creates new object, so we must clean it up
        kernel_ref = generator.get_kernel(queue_ref, kernel_module)

        ll_range = generator.create_llvm_values_for_index_space(
            ty_indexer_arg=ty_index_space,
            index_arg=ll_index_space,
        )

        device_event_ref = (
            kl.KernelLaunchIRBuilder(kernel_targetctx, builder)
            .set_kernel(kernel_ref)
            .set_queue(queue_ref)
            .set_range(
                ll_range.global_range_extents, ll_range.local_range_extents
            )
            .set_arguments(ty_kernel_args_tuple, kernel_args)
            .set_dependant_event_list()
            .submit()
        )

        # Clean up
        sycl.dpctl_kernel_delete(builder, kernel_ref)

        return wrap_event_reference(cgctx, builder, device_event_ref)

    return sig, codegen


@intrinsic(target="cpu")
def _acquire_meminfo_and_schedule_release(
    typingctx,  # pylint: disable=W0613
    ty_device_event,  # pylint: disable=W0613
    ty_kernel_args,
):
    """Generates IR code to keep arguments alive during kernel execution.

    The intrinsic collects all memory infos from the kernel arguments, acquires
    them and schecules host task to release them. Returns host task's event.
    """
    # signature of this intrinsic
    ty_event = DpctlSyclEvent()
    sig = ty_event(ty_event, ty_kernel_args)

    def codegen(cgctx, builder, sig, llargs):
        device_event = cgutils.create_struct_proxy(sig.args[0])(
            cgctx, builder, value=llargs[0]
        )

        ll_kernel_args_tuple = llargs[1]
        ty_kernel_args_tuple = sig.args[1]

        generator = CallKernelBuilder(
            codegen_targetctx=cgctx,
            builder=builder,
        )

        kernel_args = generator.extract_arguments_from_tuple(
            ty_kernel_args_tuple=ty_kernel_args_tuple,
            ll_kernel_args_tuple=ll_kernel_args_tuple,
        )

        qref = generator.get_queue_ref_val(
            kernel_argtys=ty_kernel_args_tuple,
            kernel_args=kernel_args,
        )

        host_eref = generator.acquire_meminfo_and_schedule_release(
            queue_ref=qref,
            event_ref=device_event.event_ref,
            ty_kernel_args=ty_kernel_args_tuple,
            kernel_args=kernel_args,
        )

        return wrap_event_reference(cgctx, builder, host_eref)

    return sig, codegen


@dpjit
def call_kernel(kernel_fn, index_space, *kernel_args) -> None:
    """Calls a numba_dpex.kernel decorated function from CPython or from another
    dpjit function. Kernel execution happens in syncronous way, so the thread
    will be blocked till the kernel done exectuion.

    Args:
        kernel_fn (numba_dpex.experimental.KernelDispatcher): A
        numba_dpex.kernel decorated function that is compiled to a
        KernelDispatcher by numba_dpex.
        index_space (Range | NdRange): A numba_dpex.Range or numba_dpex.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.
    """
    device_event = _submit_kernel(  # pylint: disable=E1120
        kernel_fn,
        index_space,
        kernel_args,
    )
    device_event.wait()  # pylint: disable=E1101


@dpjit
def call_kernel_async(
    kernel_fn, index_space, *kernel_args
) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]:
    """Calls a numba_dpex.kernel decorated function from CPython or from another
    dpjit function. Kernel execution happens in asyncronous way, so the thread
    will not be blocked till the kernel done exectuion. That means that it is
    user responsiblity to properly use any memory used by kernel until the
    kernel execution is completed.

    Args:
        kernel_fn (numba_dpex.experimental.KernelDispatcher): A
        numba_dpex.kernel decorated function that is compiled to a
        KernelDispatcher by numba_dpex.
        index_space (Range | NdRange): A numba_dpex.Range or numba_dpex.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:
        pair of host event and device event. Host event represent host task
        that releases use of any kernel argument so it can be deallocated.
        This task may be executed only after device task is done.
    """
    device_event = _submit_kernel(  # pylint: disable=E1120
        kernel_fn,
        index_space,
        kernel_args,
    )
    host_event = _acquire_meminfo_and_schedule_release(  # pylint: disable=E1120
        device_event, kernel_args
    )
    return host_event, device_event

1	# SPDX-FileCopyrightText: 2023 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	from typing import NamedTuple, Union	1✔
10
11	import dpctl	1✔
12	from llvmlite import ir as llvmir	1✔
13	from numba.core import cgutils, cpu, types	1✔
14	from numba.core.datamodel import default_manager as numba_default_dmm	1✔
15	from numba.core.types.containers import UniTuple	1✔
16	from numba.core.types.functions import Dispatcher	1✔
17	from numba.extending import intrinsic	1✔
18
19	from numba_dpex import config, dpjit, utils	1✔
20	from numba_dpex.core.exceptions import UnreachableError	1✔
21	from numba_dpex.core.runtime.context import DpexRTContext	1✔
22	from numba_dpex.core.types import (	1✔
23	DpctlSyclEvent,
24	DpnpNdArray,
25	NdRangeType,
26	RangeType,
27	)
28	from numba_dpex.core.utils import kernel_launcher as kl	1✔
29	from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl	1✔
30	from numba_dpex.dpctl_iface.wrappers import wrap_event_reference	1✔
31	from numba_dpex.experimental.kernel_dispatcher import (	1✔
32	KernelDispatcher,
33	_KernelModule,
34	)
35	from numba_dpex.utils import create_null_ptr	1✔
36
37
38	class LLRange(NamedTuple):	1✔
39	"""Analog of Range and NdRange but for the llvm ir values."""
40
41	global_range_extents: list	1✔
42	local_range_extents: list	1✔
43
44
45	class CallKernelBuilder:	1✔
46	"""
47	Helper class to build LLVM IR for a numba function that calls kernel.
48	"""
49
50	def __init__(	1✔
51	self,
52	codegen_targetctx: cpu.CPUContext,
53	builder: llvmir.IRBuilder,
54	):
55	self.context = codegen_targetctx	1✔
56	self.builder = builder	1✔
57	# TODO: get dpex RT from cached property once the PR is merged
58	# https://github.com/IntelPython/numba-dpex/pull/1027
59	# and get rid of the global variable. Use self.context.dpexrt instead.
60	self.dpexrt = DpexRTContext(self.context)	1✔
61
62	if config.DEBUG_KERNEL_LAUNCHER:	1!
UNCOV 63	cgutils.printf(	×
64	self.builder,
65	"DPEX-DEBUG: Inside the kernel launcher function\n",
66	)
67
68	def extract_arguments_from_tuple(	1✔
69	self,
70	ty_kernel_args_tuple: UniTuple,
71	ll_kernel_args_tuple: llvmir.Instruction,
72	) -> list[llvmir.Instruction]:
73	"""Extracts LLVM IR values from llvm tuple into python array."""
74
75	kernel_args = []	1✔
76	for pos in range(len(ty_kernel_args_tuple)):	1✔
77	kernel_args.append(	1✔
78	self.builder.extract_value(ll_kernel_args_tuple, pos)
79	)
80
81	return kernel_args	1✔
82
83	def allocate_meminfo_array(	1✔
84	self,
85	kernel_argtys: tuple[types.Type, ...],
86	kernel_args: [llvmir.Instruction, ...],
87	) -> tuple[int, list[llvmir.Instruction]]:
88	"""Allocates an LLVM array value to store each memory info from all
89	kernel arguments. The array is the populated with the LLVM value for
90	every meminfo of the kernel arguments.
91	"""
92	meminfos = []	1✔
93	for arg_num, argtype in enumerate(kernel_argtys):	1✔
94	llvm_val = kernel_args[arg_num]	1✔
95
96	meminfos += [	1✔
97	meminfo
98	for ty, meminfo in self.context.nrt.get_meminfos(
99	self.builder, argtype, llvm_val
100	)
101	]
102
103	meminfo_list = cgutils.alloca_once(	1✔
104	self.builder,
105	utils.get_llvm_type(context=self.context, type=types.voidptr),
106	size=self.context.get_constant(types.uintp, len(meminfos)),
107	)
108
109	for meminfo_num, meminfo in enumerate(meminfos):	1✔
110	meminfo_arg_dst = self.builder.gep(	1✔
111	meminfo_list,
112	[self.context.get_constant(types.int32, meminfo_num)],
113	)
114	meminfo_ptr = self.builder.bitcast(	1✔
115	meminfo,
116	utils.get_llvm_type(context=self.context, type=types.voidptr),
117	)
118	self.builder.store(meminfo_ptr, meminfo_arg_dst)	1✔
119
120	return len(meminfos), meminfo_list	1✔
121
122	def get_queue_ref_val(	1✔
123	self,
124	kernel_argtys: tuple[types.Type, ...],
125	kernel_args: [llvmir.Instruction, ...],
126	):
127	"""
128	Get the sycl queue from the first DpnpNdArray argument. Prior passes
129	before lowering make sure that compute-follows-data is enforceable
130	for a specific call to a kernel. As such, at the stage of lowering
131	the queue from the first DpnpNdArray argument can be extracted.
132	"""
133
134	for arg_num, argty in enumerate(kernel_argtys):	1!
135	if isinstance(argty, DpnpNdArray):	1!
136	llvm_val = kernel_args[arg_num]	1✔
137	datamodel = self.context.data_model_manager.lookup(argty)	1✔
138	sycl_queue_attr_pos = datamodel.get_field_position("sycl_queue")	1✔
139	ptr_to_queue_ref = self.builder.extract_value(	1✔
140	llvm_val, sycl_queue_attr_pos
141	)
142	break	1✔
143
144	return ptr_to_queue_ref	1✔
145
146	def get_kernel(self, queue_ref, kernel_module: _KernelModule):	1✔
147	"""Returns the pointer to the sycl::kernel object in a passed in
148	sycl::kernel_bundle wrapper object.
149	"""
150	# Inserts a global constant byte string in the current LLVM module to
151	# store the passed in SPIR-V binary blob.
152	kernel_bc_byte_str = self.context.insert_const_bytes(	1✔
153	self.builder.module,
154	bytes=kernel_module.kernel_bitcode,
155	)
156
157	kernel_name = self.context.insert_const_string(	1✔
158	self.builder.module, kernel_module.kernel_name
159	)
160
161	context_ref = sycl.dpctl_queue_get_context(self.builder, queue_ref)	1✔
162	device_ref = sycl.dpctl_queue_get_device(self.builder, queue_ref)	1✔
163
164	# build_or_get_kernel stills reference to context and device cause it
165	# needs to keep them alive for keys.
166	kernel_ref = self.dpexrt.build_or_get_kernel(	1✔
167	self.builder,
168	[
169	context_ref,
170	device_ref,
171	llvmir.Constant(
172	llvmir.IntType(64), hash(kernel_module.kernel_bitcode)
173	),
174	kernel_bc_byte_str,
175	llvmir.Constant(
176	llvmir.IntType(64), len(kernel_module.kernel_bitcode)
177	),
178	self.builder.load(create_null_ptr(self.builder, self.context)),
179	kernel_name,
180	],
181	)
182
183	return kernel_ref	1✔
184
185	def create_llvm_values_for_index_space(	1✔
186	self,
187	ty_indexer_arg: Union[RangeType, NdRangeType],
188	index_arg: llvmir.BaseStructType,
189	) -> LLRange:
190	"""Returns two lists of LLVM IR Values that hold the unboxed extents of
191	a Python Range or NdRange object.
192	"""
193	ndim = ty_indexer_arg.ndim	1✔
194	global_range_extents = []	1✔
195	local_range_extents = []	1✔
196	indexer_datamodel = numba_default_dmm.lookup(ty_indexer_arg)	1✔
197
198	if isinstance(ty_indexer_arg, RangeType):	1✔
199	for dim_num in range(ndim):	1✔
200	dim_pos = indexer_datamodel.get_field_position(	1✔
201	"dim" + str(dim_num)
202	)
203	global_range_extents.append(	1✔
204	self.builder.extract_value(index_arg, dim_pos)
205	)
206	elif isinstance(ty_indexer_arg, NdRangeType):	1!
207	for dim_num in range(ndim):	1✔
208	gdim_pos = indexer_datamodel.get_field_position(	1✔
209	"gdim" + str(dim_num)
210	)
211	global_range_extents.append(	1✔
212	self.builder.extract_value(index_arg, gdim_pos)
213	)
214	ldim_pos = indexer_datamodel.get_field_position(	1✔
215	"ldim" + str(dim_num)
216	)
217	local_range_extents.append(	1✔
218	self.builder.extract_value(index_arg, ldim_pos)
219	)
220	else:
UNCOV 221	raise UnreachableError	×
222
223	return LLRange(global_range_extents, local_range_extents)	1✔
224
225	def acquire_meminfo_and_schedule_release(	1✔
226	self,
227	queue_ref,
228	event_ref,
229	ty_kernel_args,
230	kernel_args,
231	):
232	"""Schedule sycl host task to release nrt meminfo of the arguments used
233	to run job. Use it to keep arguments alive during kernel execution."""
234	total_meminfos, meminfo_list = self.allocate_meminfo_array(	1✔
235	ty_kernel_args, kernel_args
236	)
237
238	event_ref_ptr = self.builder.alloca(event_ref.type)	1✔
239	self.builder.store(event_ref, event_ref_ptr)	1✔
240
241	status_ptr = cgutils.alloca_once(	1✔
242	self.builder, self.context.get_value_type(types.uint64)
243	)
244	host_eref = self.dpexrt.acquire_meminfo_and_schedule_release(	1✔
245	self.builder,
246	[
247	self.context.nrt.get_nrt_api(self.builder),
248	queue_ref,
249	meminfo_list,
250	self.context.get_constant(types.uintp, total_meminfos),
251	event_ref_ptr,
252	self.context.get_constant(types.uintp, 1),
253	status_ptr,
254	],
255	)
256	return host_eref	1✔
257
258
259	@intrinsic(target="cpu")	1✔
260	def _submit_kernel(	1✔
261	typingctx, # pylint: disable=W0613
262	ty_kernel_fn: Dispatcher,
263	ty_index_space: Union[RangeType, NdRangeType],
264	ty_kernel_args_tuple: UniTuple,
265	):
266	"""Generates IR code for call_kernel dpjit function.
267
268	The intrinsic first compiles the kernel function to SPIRV, and then to a
269	sycl kernel bundle. The arguments to the kernel are also packed into
270	flattened arrays and the sycl queue to which the kernel will be submitted
271	extracted from the args. Finally, the actual kernel is extracted from the
272	kernel bundle and submitted to the sycl queue.
273	"""
274	# signature of this intrinsic
275	ty_event = DpctlSyclEvent()	1✔
276	sig = ty_event(ty_kernel_fn, ty_index_space, ty_kernel_args_tuple)	1✔
277	kernel_sig = types.void(*ty_kernel_args_tuple)	1✔
278	# ty_kernel_fn is type specific to exact function, so we can get function
279	# directly from type and compile it. Thats why we don't need to get it in
280	# codegen
281	kernel_dispatcher: KernelDispatcher = ty_kernel_fn.dispatcher	1✔
282	kernel_dispatcher.compile(kernel_sig)	1✔
283	kernel_module: _KernelModule = kernel_dispatcher.get_overload_kcres(	1✔
284	kernel_sig
285	).kernel_device_ir_module
286	kernel_targetctx = kernel_dispatcher.targetctx	1✔
287
288	def codegen(cgctx, builder, sig, llargs):	1✔
289	# llargs[0] is kernel function that we don't need anymore (see above)
290	ty_index_space: Union[RangeType, NdRangeType] = sig.args[1]	1✔
291	ll_index_space: llvmir.Instruction = llargs[1]	1✔
292	ty_kernel_args_tuple: UniTuple = sig.args[2]	1✔
293	ll_kernel_args_tuple: llvmir.Instruction = llargs[2]	1✔
294
295	generator = CallKernelBuilder(	1✔
296	codegen_targetctx=cgctx,
297	builder=builder,
298	)
299
300	kernel_args = generator.extract_arguments_from_tuple(	1✔
301	ty_kernel_args_tuple=ty_kernel_args_tuple,
302	ll_kernel_args_tuple=ll_kernel_args_tuple,
303	)
304
305	# queue_ref is just a pointer to the attribute, so we don't have to
306	# clean it up
307	queue_ref = generator.get_queue_ref_val(	1✔
308	kernel_argtys=ty_kernel_args_tuple,
309	kernel_args=kernel_args,
310	)
311
312	# creates new object, so we must clean it up
313	kernel_ref = generator.get_kernel(queue_ref, kernel_module)	1✔
314
315	ll_range = generator.create_llvm_values_for_index_space(	1✔
316	ty_indexer_arg=ty_index_space,
317	index_arg=ll_index_space,
318	)
319
320	device_event_ref = (	1✔
321	kl.KernelLaunchIRBuilder(kernel_targetctx, builder)
322	.set_kernel(kernel_ref)
323	.set_queue(queue_ref)
324	.set_range(
325	ll_range.global_range_extents, ll_range.local_range_extents
326	)
327	.set_arguments(ty_kernel_args_tuple, kernel_args)
328	.set_dependant_event_list()
329	.submit()
330	)
331
332	# Clean up
333	sycl.dpctl_kernel_delete(builder, kernel_ref)	1✔
334
335	return wrap_event_reference(cgctx, builder, device_event_ref)	1✔
336
337	return sig, codegen	1✔
338
339
340	@intrinsic(target="cpu")	1✔
341	def _acquire_meminfo_and_schedule_release(	1✔
342	typingctx, # pylint: disable=W0613
343	ty_device_event, # pylint: disable=W0613
344	ty_kernel_args,
345	):
346	"""Generates IR code to keep arguments alive during kernel execution.
347
348	The intrinsic collects all memory infos from the kernel arguments, acquires
349	them and schecules host task to release them. Returns host task's event.
350	"""
351	# signature of this intrinsic
352	ty_event = DpctlSyclEvent()	1✔
353	sig = ty_event(ty_event, ty_kernel_args)	1✔
354
355	def codegen(cgctx, builder, sig, llargs):	1✔
356	device_event = cgutils.create_struct_proxy(sig.args[0])(	1✔
357	cgctx, builder, value=llargs[0]
358	)
359
360	ll_kernel_args_tuple = llargs[1]	1✔
361	ty_kernel_args_tuple = sig.args[1]	1✔
362
363	generator = CallKernelBuilder(	1✔
364	codegen_targetctx=cgctx,
365	builder=builder,
366	)
367
368	kernel_args = generator.extract_arguments_from_tuple(	1✔
369	ty_kernel_args_tuple=ty_kernel_args_tuple,
370	ll_kernel_args_tuple=ll_kernel_args_tuple,
371	)
372
373	qref = generator.get_queue_ref_val(	1✔
374	kernel_argtys=ty_kernel_args_tuple,
375	kernel_args=kernel_args,
376	)
377
378	host_eref = generator.acquire_meminfo_and_schedule_release(	1✔
379	queue_ref=qref,
380	event_ref=device_event.event_ref,
381	ty_kernel_args=ty_kernel_args_tuple,
382	kernel_args=kernel_args,
383	)
384
385	return wrap_event_reference(cgctx, builder, host_eref)	1✔
386
387	return sig, codegen	1✔
388
389
390	@dpjit	1✔
391	def call_kernel(kernel_fn, index_space, *kernel_args) -> None:	1✔
392	"""Calls a numba_dpex.kernel decorated function from CPython or from another
393	dpjit function. Kernel execution happens in syncronous way, so the thread
394	will be blocked till the kernel done exectuion.
395
396	Args:
397	kernel_fn (numba_dpex.experimental.KernelDispatcher): A
398	numba_dpex.kernel decorated function that is compiled to a
399	KernelDispatcher by numba_dpex.
400	index_space (Range \| NdRange): A numba_dpex.Range or numba_dpex.NdRange
401	type object that specifies the index space for the kernel.
402	kernel_args : List of objects that are passed to the numba_dpex.kernel
403	decorated function.
404	"""
UNCOV 405	device_event = _submit_kernel( # pylint: disable=E1120	×
406	kernel_fn,
407	index_space,
408	kernel_args,
409	)
UNCOV 410	device_event.wait() # pylint: disable=E1101	×
411
412
413	@dpjit	1✔
414	def call_kernel_async(	1✔
415	kernel_fn, index_space, *kernel_args
416	) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]:
417	"""Calls a numba_dpex.kernel decorated function from CPython or from another
418	dpjit function. Kernel execution happens in asyncronous way, so the thread
419	will not be blocked till the kernel done exectuion. That means that it is
420	user responsiblity to properly use any memory used by kernel until the
421	kernel execution is completed.
422
423	Args:
424	kernel_fn (numba_dpex.experimental.KernelDispatcher): A
425	numba_dpex.kernel decorated function that is compiled to a
426	KernelDispatcher by numba_dpex.
427	index_space (Range \| NdRange): A numba_dpex.Range or numba_dpex.NdRange
428	type object that specifies the index space for the kernel.
429	kernel_args : List of objects that are passed to the numba_dpex.kernel
430	decorated function.
431
432	Returns:
433	pair of host event and device event. Host event represent host task
434	that releases use of any kernel argument so it can be deallocated.
435	This task may be executed only after device task is done.
436	"""
UNCOV 437	device_event = _submit_kernel( # pylint: disable=E1120	×
438	kernel_fn,
439	index_space,
440	kernel_args,
441	)
UNCOV 442	host_event = _acquire_meminfo_and_schedule_release( # pylint: disable=E1120	×
443	device_event, kernel_args
444	)
UNCOV 445	return host_event, device_event	×

IntelPython / numba-dpex / 7199556298

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