5406637445

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

from numba.core import cgutils, types
from numba.core.ir_utils import legalize_names

from numba_dpex import utils
from numba_dpex.dpctl_iface import DpctlCAPIFnBuilder
from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum


class KernelLaunchOps:
    """
    KernelLaunchOps(lowerer, cres, num_inputs)
    Defines a set of functions to launch a SYCL kernel on the "current queue"
    as defined in the dpctl queue manager.
    """

    def _form_kernel_arg_and_arg_ty(self, val, ty):
        kernel_arg_dst = self.builder.gep(
            self.kernel_arg_array,
            [self.context.get_constant(types.int32, self.cur_arg)],
        )
        kernel_arg_ty_dst = self.builder.gep(
            self.kernel_arg_ty_array,
            [self.context.get_constant(types.int32, self.cur_arg)],
        )
        self.cur_arg += 1
        self.builder.store(val, kernel_arg_dst)
        self.builder.store(ty, kernel_arg_ty_dst)

    def __init__(self, lowerer, kernel, num_inputs):
        """Create an instance of KernelLaunchOps.

        Args:
            lowerer: The Numba Lowerer that will be used to generate the code.
            kernel: The SYCL kernel for which we are generating the code.
            num_inputs: The number of arguments to the kernels.
        """
        self.lowerer = lowerer
        self.context = self.lowerer.context
        self.builder = self.lowerer.builder
        self.kernel = kernel
        self.kernel_addr = self.kernel.addressof_ref()
        self.total_kernel_args = 0
        self.cur_arg = 0
        self.num_inputs = num_inputs
        # list of buffer that needs to comeback to host
        self.write_buffs = []
        # list of buffer that does not need to comeback to host
        self.read_only_buffs = []

    def get_current_queue(self):
        """Allocates memory on the stack to store the current queue from dpctl.

        A SYCL queue is needed to allocate USM memory and submit a kernel. This
        function gets the queue returned by ``DPCTLQueueMgr_GetCurrentQueue``
        function and stores it on the stack. The queue should be freed properly
        after returning from the kernel.

        Return: A LLVM Value storing the pointer to the SYCL queue returned
                by ``DPCTLQueueMgr_GetCurrentQueue``.

        """
        sycl_queue_val = cgutils.alloca_once(
            self.builder,
            utils.get_llvm_type(context=self.context, type=types.voidptr),
        )
        fn = DpctlCAPIFnBuilder.get_dpctl_queuemgr_get_current_queue(
            builder=self.builder, context=self.context
        )
        self.builder.store(self.builder.call(fn, []), sycl_queue_val)
        return sycl_queue_val

    def free_queue(self, sycl_queue_val):
        """
        Frees the ``DPCTLSyclQueueRef`` pointer that was used to launch the
        kernels.

        Args:
            sycl_queue_val: The SYCL queue pointer to be freed.
        """
        fn = DpctlCAPIFnBuilder.get_dpctl_queue_delete(
            builder=self.builder, context=self.context
        )
        self.builder.call(fn, [self.builder.load(sycl_queue_val)])

    def allocate_kernel_arg_array(self, num_kernel_args):
        """Allocates an array to store the LLVM Value for every kernel argument.

        Args:
            num_kernel_args : The number of kernel arguments for the kernel.

        """

        self.total_kernel_args = num_kernel_args

        # we need a kernel arg array to enqueue
        self.kernel_arg_array = cgutils.alloca_once(
            self.builder,
            utils.get_llvm_type(context=self.context, type=types.voidptr),
            size=self.context.get_constant(types.uintp, num_kernel_args),
            name="kernel_arg_array",
        )

        self.kernel_arg_ty_array = cgutils.alloca_once(
            self.builder,
            utils.LLVMTypes.int32_t,
            size=self.context.get_constant(types.uintp, num_kernel_args),
            name="kernel_arg_ty_array",
        )

    def process_kernel_arg(
        self, var, llvm_arg, arg_type, index, modified_arrays, sycl_queue_val
    ):
        """
        process_kernel_arg(var, llvm_arg, arg_type, index, modified_arrays,
        sycl_queue_val)
        Creates an LLVM Value for each kernel argument.

        Args:
            var : A kernel argument represented as a Numba type.
            llvm_arg : Only used for array arguments and points to the LLVM
            value previously allocated to store the array arg.
            arg_type : The Numba type for the argument.
            index : The poisition of the argument in the list of arguments.
            modified_arrays : The list of array arguments that are written to
            inside the kernel. The list is used to check if the argument is
            read-only or not.

        Raises:
            NotImplementedError: If an unsupported type of kernel argument is
            encountered.

        """
        if isinstance(arg_type, types.npytypes.Array):
            if llvm_arg is None:
                raise NotImplementedError(arg_type, var)

            storage = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)
            self.builder.store(
                self.context.get_constant(types.int64, 0), storage
            )
            ty = numba_type_to_dpctl_typenum(
                context=self.context, ty=types.int64
            )
            self._form_kernel_arg_and_arg_ty(
                self.builder.bitcast(
                    storage,
                    utils.get_llvm_type(
                        context=self.context, type=types.voidptr
                    ),
                ),
                ty,
            )

            storage = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)
            self.builder.store(
                self.context.get_constant(types.int64, 0), storage
            )
            ty = numba_type_to_dpctl_typenum(
                context=self.context, ty=types.int64
            )
            self._form_kernel_arg_and_arg_ty(
                self.builder.bitcast(
                    storage,
                    utils.get_llvm_type(
                        context=self.context, type=types.voidptr
                    ),
                ),
                ty,
            )

            # Handle array size
            array_size_member = self.builder.gep(
                llvm_arg,
                [
                    self.context.get_constant(types.int32, 0),
                    self.context.get_constant(types.int32, 2),
                ],
            )

            ty = numba_type_to_dpctl_typenum(
                context=self.context, ty=types.int64
            )
            self._form_kernel_arg_and_arg_ty(
                self.builder.bitcast(
                    array_size_member,
                    utils.get_llvm_type(
                        context=self.context, type=types.voidptr
                    ),
                ),
                ty,
            )

            # Handle itemsize
            item_size_member = self.builder.gep(
                llvm_arg,
                [
                    self.context.get_constant(types.int32, 0),
                    self.context.get_constant(types.int32, 3),
                ],
            )

            ty = numba_type_to_dpctl_typenum(
                context=self.context, ty=types.int64
            )
            self._form_kernel_arg_and_arg_ty(
                self.builder.bitcast(
                    item_size_member,
                    utils.get_llvm_type(
                        context=self.context, type=types.voidptr
                    ),
                ),
                ty,
            )

            # Calculate total buffer size
            total_size = cgutils.alloca_once(
                self.builder,
                utils.get_llvm_type(context=self.context, type=types.intp),
                size=utils.get_one(context=self.context),
                name="total_size" + str(self.cur_arg),
            )
            self.builder.store(
                self.builder.sext(
                    self.builder.mul(
                        self.builder.load(array_size_member),
                        self.builder.load(item_size_member),
                    ),
                    utils.get_llvm_type(context=self.context, type=types.intp),
                ),
                total_size,
            )

            # Handle data
            data_member = self.builder.gep(
                llvm_arg,
                [
                    self.context.get_constant(types.int32, 0),
                    self.context.get_constant(types.int32, 4),
                ],
            )

            # names are replaced using legalize names, we have to do the same
            # here for them to match.
            legal_names = legalize_names([var])
            ty = numba_type_to_dpctl_typenum(
                context=self.context, ty=types.voidptr
            )

            malloc_fn = DpctlCAPIFnBuilder.get_dpctl_malloc_shared(
                builder=self.builder, context=self.context
            )
            memcpy_fn = DpctlCAPIFnBuilder.get_dpctl_queue_memcpy(
                builder=self.builder, context=self.context
            )
            event_del_fn = DpctlCAPIFnBuilder.get_dpctl_event_delete(
                builder=self.builder, context=self.context
            )
            event_wait_fn = DpctlCAPIFnBuilder.get_dpctl_event_wait(
                builder=self.builder, context=self.context
            )

            # Not known to be USM so we need to copy to USM.
            buffer_name = "buffer_ptr" + str(self.cur_arg)
            # Create void * to hold new USM buffer.
            buffer_ptr = cgutils.alloca_once(
                self.builder,
                utils.get_llvm_type(context=self.context, type=types.voidptr),
                name=buffer_name,
            )
            # Setup the args to the USM allocator, size and SYCL queue.
            args = [
                self.builder.load(total_size),
                self.builder.load(sycl_queue_val),
            ]
            # Call USM shared allocator and store in buffer_ptr.
            self.builder.store(self.builder.call(malloc_fn, args), buffer_ptr)

            if legal_names[var] in modified_arrays:
                self.write_buffs.append((buffer_ptr, total_size, data_member))
            else:
                self.read_only_buffs.append(
                    (buffer_ptr, total_size, data_member)
                )

            # We really need to detect when an array needs to be copied over
            if index < self.num_inputs:
                args = [
                    self.builder.load(sycl_queue_val),
                    self.builder.load(buffer_ptr),
                    self.builder.bitcast(
                        self.builder.load(data_member),
                        utils.get_llvm_type(
                            context=self.context, type=types.voidptr
                        ),
                    ),
                    self.builder.load(total_size),
                ]
                event_ref = self.builder.call(memcpy_fn, args)
                self.builder.call(event_wait_fn, [event_ref])
                self.builder.call(event_del_fn, [event_ref])

            self._form_kernel_arg_and_arg_ty(self.builder.load(buffer_ptr), ty)

            # Handle shape
            shape_member = self.builder.gep(
                llvm_arg,
                [
                    self.context.get_constant(types.int32, 0),
                    self.context.get_constant(types.int32, 5),
                ],
            )

            for this_dim in range(arg_type.ndim):
                shape_entry = self.builder.gep(
                    shape_member,
                    [
                        self.context.get_constant(types.int32, 0),
                        self.context.get_constant(types.int32, this_dim),
                    ],
                )
                ty = numba_type_to_dpctl_typenum(
                    context=self.context, ty=types.int64
                )
                self._form_kernel_arg_and_arg_ty(
                    self.builder.bitcast(
                        shape_entry,
                        utils.get_llvm_type(
                            context=self.context, type=types.voidptr
                        ),
                    ),
                    ty,
                )

            # Handle strides
            stride_member = self.builder.gep(
                llvm_arg,
                [
                    self.context.get_constant(types.int32, 0),
                    self.context.get_constant(types.int32, 6),
                ],
            )

            for this_stride in range(arg_type.ndim):
                stride_entry = self.builder.gep(
                    stride_member,
                    [
                        self.context.get_constant(types.int32, 0),
                        self.context.get_constant(types.int32, this_stride),
                    ],
                )

                ty = numba_type_to_dpctl_typenum(
                    context=self.context, ty=types.int64
                )
                self._form_kernel_arg_and_arg_ty(
                    self.builder.bitcast(
                        stride_entry,
                        utils.get_llvm_type(
                            context=self.context, type=types.voidptr
                        ),
                    ),
                    ty,
                )

        else:
            ty = numba_type_to_dpctl_typenum(context=self.context, ty=arg_type)
            self._form_kernel_arg_and_arg_ty(
                self.builder.bitcast(
                    llvm_arg,
                    utils.get_llvm_type(
                        context=self.context, type=types.voidptr
                    ),
                ),
                ty,
            )

    def enqueue_kernel_and_copy_back(self, dim_bounds, sycl_queue_val):
        """
        enqueue_kernel_and_copy_back(dim_bounds, sycl_queue_val)
        Submits the kernel to the specified queue, waits and then copies
        back any results to the host.

        Args:
            dim_bounds : An array of three tuple representing the starting
                         offset, end offset and the stride (step) for each
                         dimension of the input arrays. Every array in a parfor
                         is of the same dimensionality and shape, thus ensuring
                         the bounds are the same.
            sycl_queue_val : The SYCL queue on which the kernel is
                             submitted.
        """
        submit_fn = DpctlCAPIFnBuilder.get_dpctl_queue_submit_range(
            builder=self.builder, context=self.context
        )
        queue_wait_fn = DpctlCAPIFnBuilder.get_dpctl_queue_wait(
            builder=self.builder, context=self.context
        )
        event_del_fn = DpctlCAPIFnBuilder.get_dpctl_event_delete(
            builder=self.builder, context=self.context
        )
        memcpy_fn = DpctlCAPIFnBuilder.get_dpctl_queue_memcpy(
            builder=self.builder, context=self.context
        )
        free_fn = DpctlCAPIFnBuilder.get_dpctl_free_with_queue(
            builder=self.builder, context=self.context
        )
        event_wait_fn = DpctlCAPIFnBuilder.get_dpctl_event_wait(
            builder=self.builder, context=self.context
        )

        # the assumption is loop_ranges will always be less than or equal to 3
        # dimensions
        num_dim = len(dim_bounds) if len(dim_bounds) < 4 else 3

        # form the global range
        global_range = cgutils.alloca_once(
            self.builder,
            utils.get_llvm_type(context=self.context, type=types.uintp),
            size=self.context.get_constant(types.uintp, num_dim),
            name="global_range",
        )

        intp_t = utils.get_llvm_type(context=self.context, type=types.intp)
        intp_ptr_t = utils.get_llvm_ptr_type(intp_t)

        for i in range(num_dim):
            start, stop, step = dim_bounds[i]
            if stop.type != utils.LLVMTypes.int64_t:
                stop = self.builder.sext(stop, utils.LLVMTypes.int64_t)

            # we reverse the global range to account for how sycl and opencl
            # range differs
            self.builder.store(
                stop,
                self.builder.gep(
                    global_range,
                    [self.context.get_constant(types.uintp, (num_dim - 1) - i)],
                ),
            )

        args = [
            self.builder.inttoptr(
                self.context.get_constant(types.uintp, self.kernel_addr),
                utils.get_llvm_type(context=self.context, type=types.voidptr),
            ),
            self.builder.load(sycl_queue_val),
            self.kernel_arg_array,
            self.kernel_arg_ty_array,
            self.context.get_constant(types.uintp, self.total_kernel_args),
            self.builder.bitcast(global_range, intp_ptr_t),
            self.context.get_constant(types.uintp, num_dim),
            self.builder.bitcast(
                utils.create_null_ptr(
                    builder=self.builder, context=self.context
                ),
                utils.get_llvm_type(context=self.context, type=types.voidptr),
            ),
            self.context.get_constant(types.uintp, 0),
        ]

        # Submit the kernel
        event_ref = self.builder.call(submit_fn, args)

        # Add a wait on the queue
        self.builder.call(queue_wait_fn, [self.builder.load(sycl_queue_val)])

        # Note that the dpctl_queue_wait call waits on the event and then
        # decrements the ref count of the sycl::event C++ object. However, the
        # event object returned by the get_dpctl_queue_submit_range call still
        # needs to be explicitly deleted to free up the event object properly.
        self.builder.call(event_del_fn, [event_ref])

        # read buffers back to host
        for write_buff in self.write_buffs:
            buffer_ptr, total_size, data_member = write_buff
            args = [
                self.builder.load(sycl_queue_val),
                self.builder.bitcast(
                    self.builder.load(data_member),
                    utils.get_llvm_type(
                        context=self.context, type=types.voidptr
                    ),
                ),
                self.builder.load(buffer_ptr),
                self.builder.load(total_size),
            ]
            # FIXME: In future, when the DctlQueue_Memcpy is made non-blocking
            # the returned event should be explicitly freed by calling
            # get_dpctl_event_delete.
            event_ref = self.builder.call(memcpy_fn, args)
            self.builder.call(event_wait_fn, [event_ref])
            self.builder.call(event_del_fn, [event_ref])

            self.builder.call(
                free_fn,
                [
                    self.builder.load(buffer_ptr),
                    self.builder.load(sycl_queue_val),
                ],
            )

        for read_buff in self.read_only_buffs:
            buffer_ptr, total_size, data_member = read_buff
            self.builder.call(
                free_fn,
                [
                    self.builder.load(buffer_ptr),
                    self.builder.load(sycl_queue_val),
                ],
            )

1	# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	from numba.core import cgutils, types	1✔
6	from numba.core.ir_utils import legalize_names	1✔
7
8	from numba_dpex import utils	1✔
9	from numba_dpex.dpctl_iface import DpctlCAPIFnBuilder	1✔
10	from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum	1✔
11
12
13	class KernelLaunchOps:	1✔
14	"""
15	KernelLaunchOps(lowerer, cres, num_inputs)
16	Defines a set of functions to launch a SYCL kernel on the "current queue"
17	as defined in the dpctl queue manager.
18	"""
19
20	def _form_kernel_arg_and_arg_ty(self, val, ty):	1✔
21	kernel_arg_dst = self.builder.gep(	×
22	self.kernel_arg_array,
23	[self.context.get_constant(types.int32, self.cur_arg)],
24	)
25	kernel_arg_ty_dst = self.builder.gep(	×
26	self.kernel_arg_ty_array,
27	[self.context.get_constant(types.int32, self.cur_arg)],
28	)
29	self.cur_arg += 1	×
30	self.builder.store(val, kernel_arg_dst)	×
31	self.builder.store(ty, kernel_arg_ty_dst)	×
32
33	def __init__(self, lowerer, kernel, num_inputs):	1✔
34	"""Create an instance of KernelLaunchOps.
35
36	Args:
37	lowerer: The Numba Lowerer that will be used to generate the code.
38	kernel: The SYCL kernel for which we are generating the code.
39	num_inputs: The number of arguments to the kernels.
40	"""
41	self.lowerer = lowerer	×
42	self.context = self.lowerer.context	×
43	self.builder = self.lowerer.builder	×
44	self.kernel = kernel	×
45	self.kernel_addr = self.kernel.addressof_ref()	×
46	self.total_kernel_args = 0	×
47	self.cur_arg = 0	×
48	self.num_inputs = num_inputs	×
49	# list of buffer that needs to comeback to host
50	self.write_buffs = []	×
51	# list of buffer that does not need to comeback to host
52	self.read_only_buffs = []	×
53
54	def get_current_queue(self):	1✔
55	"""Allocates memory on the stack to store the current queue from dpctl.
56
57	A SYCL queue is needed to allocate USM memory and submit a kernel. This
58	function gets the queue returned by ``DPCTLQueueMgr_GetCurrentQueue``
59	function and stores it on the stack. The queue should be freed properly
60	after returning from the kernel.
61
62	Return: A LLVM Value storing the pointer to the SYCL queue returned
63	by ``DPCTLQueueMgr_GetCurrentQueue``.
64
65	"""
66	sycl_queue_val = cgutils.alloca_once(	×
67	self.builder,
68	utils.get_llvm_type(context=self.context, type=types.voidptr),
69	)
70	fn = DpctlCAPIFnBuilder.get_dpctl_queuemgr_get_current_queue(	×
71	builder=self.builder, context=self.context
72	)
73	self.builder.store(self.builder.call(fn, []), sycl_queue_val)	×
74	return sycl_queue_val	×
75
76	def free_queue(self, sycl_queue_val):	1✔
77	"""
78	Frees the ``DPCTLSyclQueueRef`` pointer that was used to launch the
79	kernels.
80
81	Args:
82	sycl_queue_val: The SYCL queue pointer to be freed.
83	"""
84	fn = DpctlCAPIFnBuilder.get_dpctl_queue_delete(	×
85	builder=self.builder, context=self.context
86	)
87	self.builder.call(fn, [self.builder.load(sycl_queue_val)])	×
88
89	def allocate_kernel_arg_array(self, num_kernel_args):	1✔
90	"""Allocates an array to store the LLVM Value for every kernel argument.
91
92	Args:
93	num_kernel_args : The number of kernel arguments for the kernel.
94
95	"""
96
97	self.total_kernel_args = num_kernel_args	×
98
99	# we need a kernel arg array to enqueue
100	self.kernel_arg_array = cgutils.alloca_once(	×
101	self.builder,
102	utils.get_llvm_type(context=self.context, type=types.voidptr),
103	size=self.context.get_constant(types.uintp, num_kernel_args),
104	name="kernel_arg_array",
105	)
106
107	self.kernel_arg_ty_array = cgutils.alloca_once(	×
108	self.builder,
109	utils.LLVMTypes.int32_t,
110	size=self.context.get_constant(types.uintp, num_kernel_args),
111	name="kernel_arg_ty_array",
112	)
113
114	def process_kernel_arg(	1✔
115	self, var, llvm_arg, arg_type, index, modified_arrays, sycl_queue_val
116	):
117	"""
118	process_kernel_arg(var, llvm_arg, arg_type, index, modified_arrays,
119	sycl_queue_val)
120	Creates an LLVM Value for each kernel argument.
121
122	Args:
123	var : A kernel argument represented as a Numba type.
124	llvm_arg : Only used for array arguments and points to the LLVM
125	value previously allocated to store the array arg.
126	arg_type : The Numba type for the argument.
127	index : The poisition of the argument in the list of arguments.
128	modified_arrays : The list of array arguments that are written to
129	inside the kernel. The list is used to check if the argument is
130	read-only or not.
131
132	Raises:
133	NotImplementedError: If an unsupported type of kernel argument is
134	encountered.
135
136	"""
137	if isinstance(arg_type, types.npytypes.Array):	×
138	if llvm_arg is None:	×
139	raise NotImplementedError(arg_type, var)
140
141	storage = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)	×
142	self.builder.store(	×
143	self.context.get_constant(types.int64, 0), storage
144	)
145	ty = numba_type_to_dpctl_typenum(	×
146	context=self.context, ty=types.int64
147	)
148	self._form_kernel_arg_and_arg_ty(	×
149	self.builder.bitcast(
150	storage,
151	utils.get_llvm_type(
152	context=self.context, type=types.voidptr
153	),
154	),
155	ty,
156	)
157
158	storage = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)	×
159	self.builder.store(	×
160	self.context.get_constant(types.int64, 0), storage
161	)
162	ty = numba_type_to_dpctl_typenum(	×
163	context=self.context, ty=types.int64
164	)
165	self._form_kernel_arg_and_arg_ty(	×
166	self.builder.bitcast(
167	storage,
168	utils.get_llvm_type(
169	context=self.context, type=types.voidptr
170	),
171	),
172	ty,
173	)
174
175	# Handle array size
176	array_size_member = self.builder.gep(	×
177	llvm_arg,
178	[
179	self.context.get_constant(types.int32, 0),
180	self.context.get_constant(types.int32, 2),
181	],
182	)
183
184	ty = numba_type_to_dpctl_typenum(	×
185	context=self.context, ty=types.int64
186	)
187	self._form_kernel_arg_and_arg_ty(	×
188	self.builder.bitcast(
189	array_size_member,
190	utils.get_llvm_type(
191	context=self.context, type=types.voidptr
192	),
193	),
194	ty,
195	)
196
197	# Handle itemsize
198	item_size_member = self.builder.gep(	×
199	llvm_arg,
200	[
201	self.context.get_constant(types.int32, 0),
202	self.context.get_constant(types.int32, 3),
203	],
204	)
205
206	ty = numba_type_to_dpctl_typenum(	×
207	context=self.context, ty=types.int64
208	)
209	self._form_kernel_arg_and_arg_ty(	×
210	self.builder.bitcast(
211	item_size_member,
212	utils.get_llvm_type(
213	context=self.context, type=types.voidptr
214	),
215	),
216	ty,
217	)
218
219	# Calculate total buffer size
220	total_size = cgutils.alloca_once(	×
221	self.builder,
222	utils.get_llvm_type(context=self.context, type=types.intp),
223	size=utils.get_one(context=self.context),
224	name="total_size" + str(self.cur_arg),
225	)
226	self.builder.store(	×
227	self.builder.sext(
228	self.builder.mul(
229	self.builder.load(array_size_member),
230	self.builder.load(item_size_member),
231	),
232	utils.get_llvm_type(context=self.context, type=types.intp),
233	),
234	total_size,
235	)
236
237	# Handle data
238	data_member = self.builder.gep(	×
239	llvm_arg,
240	[
241	self.context.get_constant(types.int32, 0),
242	self.context.get_constant(types.int32, 4),
243	],
244	)
245
246	# names are replaced using legalize names, we have to do the same
247	# here for them to match.
248	legal_names = legalize_names([var])	×
249	ty = numba_type_to_dpctl_typenum(	×
250	context=self.context, ty=types.voidptr
251	)
252
253	malloc_fn = DpctlCAPIFnBuilder.get_dpctl_malloc_shared(	×
254	builder=self.builder, context=self.context
255	)
256	memcpy_fn = DpctlCAPIFnBuilder.get_dpctl_queue_memcpy(	×
257	builder=self.builder, context=self.context
258	)
259	event_del_fn = DpctlCAPIFnBuilder.get_dpctl_event_delete(	×
260	builder=self.builder, context=self.context
261	)
262	event_wait_fn = DpctlCAPIFnBuilder.get_dpctl_event_wait(	×
263	builder=self.builder, context=self.context
264	)
265
266	# Not known to be USM so we need to copy to USM.
267	buffer_name = "buffer_ptr" + str(self.cur_arg)	×
268	# Create void * to hold new USM buffer.
269	buffer_ptr = cgutils.alloca_once(	×
270	self.builder,
271	utils.get_llvm_type(context=self.context, type=types.voidptr),
272	name=buffer_name,
273	)
274	# Setup the args to the USM allocator, size and SYCL queue.
275	args = [	×
276	self.builder.load(total_size),
277	self.builder.load(sycl_queue_val),
278	]
279	# Call USM shared allocator and store in buffer_ptr.
280	self.builder.store(self.builder.call(malloc_fn, args), buffer_ptr)	×
281
282	if legal_names[var] in modified_arrays:	×
283	self.write_buffs.append((buffer_ptr, total_size, data_member))	×
284	else:
285	self.read_only_buffs.append(	×
286	(buffer_ptr, total_size, data_member)
287	)
288
289	# We really need to detect when an array needs to be copied over
290	if index < self.num_inputs:	×
291	args = [	×
292	self.builder.load(sycl_queue_val),
293	self.builder.load(buffer_ptr),
294	self.builder.bitcast(
295	self.builder.load(data_member),
296	utils.get_llvm_type(
297	context=self.context, type=types.voidptr
298	),
299	),
300	self.builder.load(total_size),
301	]
302	event_ref = self.builder.call(memcpy_fn, args)	×
303	self.builder.call(event_wait_fn, [event_ref])	×
304	self.builder.call(event_del_fn, [event_ref])	×
305
306	self._form_kernel_arg_and_arg_ty(self.builder.load(buffer_ptr), ty)	×
307
308	# Handle shape
309	shape_member = self.builder.gep(	×
310	llvm_arg,
311	[
312	self.context.get_constant(types.int32, 0),
313	self.context.get_constant(types.int32, 5),
314	],
315	)
316
317	for this_dim in range(arg_type.ndim):	×
318	shape_entry = self.builder.gep(	×
319	shape_member,
320	[
321	self.context.get_constant(types.int32, 0),
322	self.context.get_constant(types.int32, this_dim),
323	],
324	)
325	ty = numba_type_to_dpctl_typenum(	×
326	context=self.context, ty=types.int64
327	)
328	self._form_kernel_arg_and_arg_ty(	×
329	self.builder.bitcast(
330	shape_entry,
331	utils.get_llvm_type(
332	context=self.context, type=types.voidptr
333	),
334	),
335	ty,
336	)
337
338	# Handle strides
339	stride_member = self.builder.gep(	×
340	llvm_arg,
341	[
342	self.context.get_constant(types.int32, 0),
343	self.context.get_constant(types.int32, 6),
344	],
345	)
346
347	for this_stride in range(arg_type.ndim):	×
348	stride_entry = self.builder.gep(	×
349	stride_member,
350	[
351	self.context.get_constant(types.int32, 0),
352	self.context.get_constant(types.int32, this_stride),
353	],
354	)
355
356	ty = numba_type_to_dpctl_typenum(	×
357	context=self.context, ty=types.int64
358	)
359	self._form_kernel_arg_and_arg_ty(	×
360	self.builder.bitcast(
361	stride_entry,
362	utils.get_llvm_type(
363	context=self.context, type=types.voidptr
364	),
365	),
366	ty,
367	)
368
369	else:
NEW 370	ty = numba_type_to_dpctl_typenum(context=self.context, ty=arg_type)	×
UNCOV 371	self._form_kernel_arg_and_arg_ty(	×
372	self.builder.bitcast(
373	llvm_arg,
374	utils.get_llvm_type(
375	context=self.context, type=types.voidptr
376	),
377	),
378	ty,
379	)
380
381	def enqueue_kernel_and_copy_back(self, dim_bounds, sycl_queue_val):	1✔
382	"""
383	enqueue_kernel_and_copy_back(dim_bounds, sycl_queue_val)
384	Submits the kernel to the specified queue, waits and then copies
385	back any results to the host.
386
387	Args:
388	dim_bounds : An array of three tuple representing the starting
389	offset, end offset and the stride (step) for each
390	dimension of the input arrays. Every array in a parfor
391	is of the same dimensionality and shape, thus ensuring
392	the bounds are the same.
393	sycl_queue_val : The SYCL queue on which the kernel is
394	submitted.
395	"""
396	submit_fn = DpctlCAPIFnBuilder.get_dpctl_queue_submit_range(	×
397	builder=self.builder, context=self.context
398	)
399	queue_wait_fn = DpctlCAPIFnBuilder.get_dpctl_queue_wait(	×
400	builder=self.builder, context=self.context
401	)
402	event_del_fn = DpctlCAPIFnBuilder.get_dpctl_event_delete(	×
403	builder=self.builder, context=self.context
404	)
405	memcpy_fn = DpctlCAPIFnBuilder.get_dpctl_queue_memcpy(	×
406	builder=self.builder, context=self.context
407	)
408	free_fn = DpctlCAPIFnBuilder.get_dpctl_free_with_queue(	×
409	builder=self.builder, context=self.context
410	)
411	event_wait_fn = DpctlCAPIFnBuilder.get_dpctl_event_wait(	×
412	builder=self.builder, context=self.context
413	)
414
415	# the assumption is loop_ranges will always be less than or equal to 3
416	# dimensions
417	num_dim = len(dim_bounds) if len(dim_bounds) < 4 else 3	×
418
419	# form the global range
420	global_range = cgutils.alloca_once(	×
421	self.builder,
422	utils.get_llvm_type(context=self.context, type=types.uintp),
423	size=self.context.get_constant(types.uintp, num_dim),
424	name="global_range",
425	)
426
427	intp_t = utils.get_llvm_type(context=self.context, type=types.intp)	×
428	intp_ptr_t = utils.get_llvm_ptr_type(intp_t)	×
429
430	for i in range(num_dim):	×
431	start, stop, step = dim_bounds[i]	×
432	if stop.type != utils.LLVMTypes.int64_t:	×
433	stop = self.builder.sext(stop, utils.LLVMTypes.int64_t)	×
434
435	# we reverse the global range to account for how sycl and opencl
436	# range differs
437	self.builder.store(	×
438	stop,
439	self.builder.gep(
440	global_range,
441	[self.context.get_constant(types.uintp, (num_dim - 1) - i)],
442	),
443	)
444
445	args = [	×
446	self.builder.inttoptr(
447	self.context.get_constant(types.uintp, self.kernel_addr),
448	utils.get_llvm_type(context=self.context, type=types.voidptr),
449	),
450	self.builder.load(sycl_queue_val),
451	self.kernel_arg_array,
452	self.kernel_arg_ty_array,
453	self.context.get_constant(types.uintp, self.total_kernel_args),
454	self.builder.bitcast(global_range, intp_ptr_t),
455	self.context.get_constant(types.uintp, num_dim),
456	self.builder.bitcast(
457	utils.create_null_ptr(
458	builder=self.builder, context=self.context
459	),
460	utils.get_llvm_type(context=self.context, type=types.voidptr),
461	),
462	self.context.get_constant(types.uintp, 0),
463	]
464
465	# Submit the kernel
466	event_ref = self.builder.call(submit_fn, args)	×
467
468	# Add a wait on the queue
469	self.builder.call(queue_wait_fn, [self.builder.load(sycl_queue_val)])	×
470
471	# Note that the dpctl_queue_wait call waits on the event and then
472	# decrements the ref count of the sycl::event C++ object. However, the
473	# event object returned by the get_dpctl_queue_submit_range call still
474	# needs to be explicitly deleted to free up the event object properly.
475	self.builder.call(event_del_fn, [event_ref])	×
476
477	# read buffers back to host
478	for write_buff in self.write_buffs:	×
479	buffer_ptr, total_size, data_member = write_buff	×
480	args = [	×
481	self.builder.load(sycl_queue_val),
482	self.builder.bitcast(
483	self.builder.load(data_member),
484	utils.get_llvm_type(
485	context=self.context, type=types.voidptr
486	),
487	),
488	self.builder.load(buffer_ptr),
489	self.builder.load(total_size),
490	]
491	# FIXME: In future, when the DctlQueue_Memcpy is made non-blocking
492	# the returned event should be explicitly freed by calling
493	# get_dpctl_event_delete.
494	event_ref = self.builder.call(memcpy_fn, args)	×
495	self.builder.call(event_wait_fn, [event_ref])	×
496	self.builder.call(event_del_fn, [event_ref])	×
497
498	self.builder.call(	×
499	free_fn,
500	[
501	self.builder.load(buffer_ptr),
502	self.builder.load(sycl_queue_val),
503	],
504	)
505
506	for read_buff in self.read_only_buffs:	×
507	buffer_ptr, total_size, data_member = read_buff	×
508	self.builder.call(	×
509	free_fn,
510	[
511	self.builder.load(buffer_ptr),
512	self.builder.load(sycl_queue_val),
513	],
514	)

IntelPython / numba-dpex / 5406637445

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