7201235944

Committed 13 Dec 2023 09:25PM UTC coverage: 79.45% (+0.1%) from 79.335%

Build # 7201235944

Build Type

Pull #1243

github

Committed by

web-flow

Commit Message

Merge 439b5f9ba into f4b618f66

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

Run Details

1195 of 1762 branches covered (0.0%)

Branch coverage included in aggregate %.

115 of 116 new or added lines in 5 files covered. (99.14%)

24 existing lines in 4 files now uncovered.

5335 of 6457 relevant lines covered (82.62%)

0.83 hits per line

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

97.03

/numba_dpex/core/utils/kernel_launcher.py

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

"""Module that contains numba style wrapper around sycl kernel submit."""

from dataclasses import dataclass

import dpctl
from llvmlite import ir as llvmir
from llvmlite.ir.builder import IRBuilder
from numba.core import cgutils, types

from numba_dpex import config, utils
from numba_dpex.core.runtime.context import DpexRTContext
from numba_dpex.core.types import DpnpNdArray
from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum

MAX_SIZE_OF_SYCL_RANGE = 3


@dataclass
class _KernelLaunchIRArguments:  # pylint: disable=too-many-instance-attributes
    """List of kernel launch arguments used in sycl.dpctl_queue_submit_range and
    sycl.dpctl_queue_submit_ndrange."""

    sycl_kernel_ref: llvmir.Instruction = None
    sycl_queue_ref: llvmir.Instruction = None
    arg_list: llvmir.Instruction = None
    arg_ty_list: llvmir.Instruction = None
    total_kernel_args: llvmir.Instruction = None
    global_range: llvmir.Instruction = None
    local_range: llvmir.Instruction = None
    range_size: llvmir.Instruction = None
    dep_events: llvmir.Instruction = None
    dep_events_len: llvmir.Instruction = None

    def to_list(self):
        """Returns list of arguments in the right order to pass to
        sycl.dpctl_queue_submit_range or sycl.dpctl_queue_submit_ndrange."""
        res = [
            self.sycl_kernel_ref,
            self.sycl_queue_ref,
            self.arg_list,
            self.arg_ty_list,
            self.total_kernel_args,
            self.global_range,
        ]

        if self.local_range is not None:
            res += [self.local_range]

        res += [
            self.range_size,
            self.dep_events,
            self.dep_events_len,
        ]

        return res


class KernelLaunchIRBuilder:
    """
    KernelLaunchIRBuilder(lowerer, cres)
    Helper class to build the LLVM IR for the submission of a kernel.

    The class generates LLVM IR inside the current LLVM module that is needed
    for submitting kernels. The LLVM Values that
    """

    def __init__(self, context, builder: IRBuilder):
        """Create a KernelLauncher for the specified kernel.

        Args:
            context: A Numba target context that will be used to generate the
                     code.
            builder: An llvmlite IRBuilder instance used to generate LLVM IR.
        """
        self.context = context
        self.builder = builder
        self.rtctx = DpexRTContext(self.context)
        self.arguments = _KernelLaunchIRArguments()

    def _build_nullptr(self):
        """Builds the LLVM IR to represent a null pointer.

        Returns: An LLVM Value storing a null pointer
        """
        zero = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)
        self.builder.store(self.context.get_constant(types.int64, 0), zero)
        return self.builder.bitcast(
            zero, utils.get_llvm_type(context=self.context, type=types.voidptr)
        )

    def _build_array_attr_arg(  # pylint: disable=too-many-arguments
        self,
        array_val,
        array_attr_pos,
        array_attr_ty,
        arg_list,
        args_ty_list,
        arg_num,
    ):
        array_attr = self.builder.gep(
            array_val,
            [
                self.context.get_constant(types.int32, 0),
                self.context.get_constant(types.int32, array_attr_pos),
            ],
        )

        # FIXME: If pointer arg then load it to some value and pass that value.
        # We also most likely need an address space cast
        if isinstance(
            array_attr_ty, (types.misc.RawPointer, types.misc.CPointer)
        ):
            array_attr = self.builder.load(array_attr)

        self._build_arg(
            val=array_attr,
            ty=array_attr_ty,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )

    def _build_unituple_member_arg(  # pylint: disable=too-many-arguments
        self, array_val, array_attr_pos, ndims, arg_list, args_ty_list, arg_num
    ):
        array_attr = self.builder.gep(
            array_val,
            [
                self.context.get_constant(types.int32, 0),
                self.context.get_constant(types.int32, array_attr_pos),
            ],
        )

        for ndim in range(ndims):
            self._build_array_attr_arg(
                array_val=array_attr,
                array_attr_pos=ndim,
                array_attr_ty=types.int64,
                arg_list=arg_list,
                args_ty_list=args_ty_list,
                arg_num=arg_num + ndim,
            )

    def _build_arg(
        self, val, ty, arg_list, args_ty_list, arg_num
    ):  # pylint: disable=too-many-arguments
        """Stores the kernel arguments and the kernel argument types into
        arrays that will be passed to DPCTLQueue_SubmitRange.

        Args:
            val: An LLVM IR Value that will be stored into the arguments array
            ty: A Numba type that will be converted to a DPCTLKernelArgType
            enum and stored into the argument types list array
            arg_list: An LLVM IR Value array that stores the kernel arguments
            args_ty_list: An LLVM IR Value array that stores the
            DPCTLKernelArgType enum value for each kernel argument
            arg_num: The index position at which the arg_list and args_ty_list
            need to be updated.
        """
        kernel_arg_dst = self.builder.gep(
            arg_list,
            [self.context.get_constant(types.int32, arg_num)],
        )
        kernel_arg_ty_dst = self.builder.gep(
            args_ty_list,
            [self.context.get_constant(types.int32, arg_num)],
        )
        val = self.builder.bitcast(
            val,
            utils.get_llvm_type(context=self.context, type=types.voidptr),
        )
        self.builder.store(val, kernel_arg_dst)
        self.builder.store(
            numba_type_to_dpctl_typenum(self.context, ty), kernel_arg_ty_dst
        )

    def _build_complex_arg(
        self, val, ty, arg_list, args_ty_list, arg_num
    ):  # pylint: disable=too-many-arguments
        """Creates a list of LLVM Values for an unpacked complex kernel
        argument.
        """
        self._build_array_attr_arg(
            array_val=val,
            array_attr_pos=0,
            array_attr_ty=ty,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        self._build_array_attr_arg(
            array_val=val,
            array_attr_pos=1,
            array_attr_ty=ty,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1

    def _build_array_arg(  # pylint: disable=too-many-arguments
        self,
        array_val,
        array_data_model,
        array_rank,  # pylint: disable=unused-argument
        arg_list,
        args_ty_list,
        arg_num,
    ):
        """Creates a list of LLVM Values for an unpacked DpnpNdArray kernel
        argument.

        The steps performed here are the same as in
        numba_dpex.core.kernel_interface.arg_pack_unpacker._unpack_array_helper
        """
        # Argument 1: Null pointer for the NRT_MemInfo attribute of the array
        nullptr = self._build_nullptr()
        self._build_arg(
            val=nullptr,
            ty=types.int64,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        # Argument 2: Null pointer for the Parent attribute of the array
        nullptr = self._build_nullptr()
        self._build_arg(
            val=nullptr,
            ty=types.int64,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        # Argument nitems
        self._build_array_attr_arg(
            array_val=array_val,
            array_attr_pos=array_data_model.get_field_position("nitems"),
            array_attr_ty=array_data_model.get_member_fe_type("nitems"),
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        # Argument itemsize
        self._build_array_attr_arg(
            array_val=array_val,
            array_attr_pos=array_data_model.get_field_position("itemsize"),
            array_attr_ty=array_data_model.get_member_fe_type("itemsize"),
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        # Argument data
        self._build_array_attr_arg(
            array_val=array_val,
            array_attr_pos=array_data_model.get_field_position("data"),
            array_attr_ty=array_data_model.get_member_fe_type("data"),
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        # Argument sycl_queue: as the queue pointer is not to be used in a
        # kernel we always pass in a nullptr
        self._build_arg(
            val=nullptr,
            ty=types.int64,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += 1
        # Arguments for shape
        shape_member = array_data_model.get_member_fe_type("shape")
        self._build_unituple_member_arg(
            array_val=array_val,
            array_attr_pos=array_data_model.get_field_position("shape"),
            ndims=shape_member.count,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += shape_member.count
        # Arguments for strides
        stride_member = array_data_model.get_member_fe_type("strides")
        self._build_unituple_member_arg(
            array_val=array_val,
            array_attr_pos=array_data_model.get_field_position("strides"),
            ndims=stride_member.count,
            arg_list=arg_list,
            args_ty_list=args_ty_list,
            arg_num=arg_num,
        )
        arg_num += stride_member.count

    # TODO: remove, not part of the builder
    def get_queue(self, exec_queue: dpctl.SyclQueue) -> llvmir.Instruction:
        """Allocates memory on the stack to store a DPCTLSyclQueueRef.

        Returns: A LLVM Value storing the pointer to the SYCL queue created
        using the filter string for the Python exec_queue (dpctl.SyclQueue).
        """

        # Allocate a stack var to store the queue created from the filter string
        sycl_queue_val = cgutils.alloca_once(
            self.builder,
            utils.get_llvm_type(context=self.context, type=types.voidptr),
        )
        # Insert a global constant to store the filter string
        device = self.context.insert_const_string(
            self.builder.module, exec_queue.sycl_device.filter_string
        )
        # Store the queue returned by DPEXRTQueue_CreateFromFilterString in a
        # local variable
        self.builder.store(
            self.rtctx.get_queue_from_filter_string(
                builder=self.builder, device=device
            ),
            sycl_queue_val,
        )
        return 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 (int): The number of kernel arguments that
            determines the size of args array to allocate.

        Returns: An LLVM IR value pointing to an array to store the kernel
        arguments.
        """
        args_list = cgutils.alloca_once(
            self.builder,
            utils.LLVMTypes.byte_ptr_t,
            size=self.context.get_constant(types.uintp, num_kernel_args),
        )

        return args_list

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

        Args:
            num_kernel_args (int): The number of kernel arguments that
            determines the size of args array to allocate.

        Returns: An LLVM IR value pointing to an array to store the kernel
        arguments typenums as defined in dpctl.
        """
        args_ty_list = cgutils.alloca_once(
            self.builder,
            utils.LLVMTypes.int32_t,
            size=self.context.get_constant(types.uintp, num_kernel_args),
        )

        return args_ty_list

    def _create_sycl_range(self, idx_range):
        """Allocate a size_t[3] array to store the extents of a sycl::range.

        Sycl supports upto 3-dimensional ranges and a such the array is
        statically sized to length three. Only the elements that store an actual
        range value are populated based on the size of the idx_range argument.

        """
        intp_t = utils.get_llvm_type(context=self.context, type=types.intp)
        intp_ptr_t = utils.get_llvm_ptr_type(intp_t)
        num_dim = len(idx_range)

        # form the global range
        range_list = cgutils.alloca_once(
            self.builder,
            utils.get_llvm_type(context=self.context, type=types.uintp),
            size=self.context.get_constant(types.uintp, MAX_SIZE_OF_SYCL_RANGE),
        )

        for i in range(num_dim):
            rext = idx_range[i]
            if rext.type != utils.LLVMTypes.int64_t:
                rext = self.builder.sext(rext, utils.LLVMTypes.int64_t)

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

        return self.builder.bitcast(range_list, intp_ptr_t)

    def set_kernel(self, sycl_kernel_ref: llvmir.Instruction):
        """Sets kernel to the argument list."""
        self.arguments.sycl_kernel_ref = sycl_kernel_ref

    def set_queue(self, sycl_queue_ref: llvmir.Instruction):
        """Sets queue to the argument list."""
        self.arguments.sycl_queue_ref = sycl_queue_ref

    def set_range(
        self,
        global_range: list,
        local_range: list = None,
    ):
        """Sets global and local range if provided to the argument list."""
        self.arguments.global_range = self._create_sycl_range(global_range)
        if local_range is not None and len(local_range) > 0:
            self.arguments.local_range = self._create_sycl_range(local_range)
        self.arguments.range_size = self.context.get_constant(
            types.uintp, len(global_range)
        )

    def set_arguments(
        self,
        ty_kernel_args: list,
        kernel_args: list,
    ):
        """Sets flattened kernel args, kernel arg types and number of those
        arguments to the argument list."""
        if config.DEBUG_KERNEL_LAUNCHER:
            cgutils.printf(
                self.builder,
                "DPEX-DEBUG: Populating kernel args and arg type arrays.\n",
            )

        num_flattened_kernel_args = self._get_num_flattened_kernel_args(
            kernel_argtys=ty_kernel_args,
        )

        # Create LLVM values for the kernel args list and kernel arg types list
        args_list = self._allocate_kernel_arg_array(num_flattened_kernel_args)

        args_ty_list = self._allocate_kernel_arg_ty_array(
            num_flattened_kernel_args
        )

        kernel_args_ptrs = []
        for arg in kernel_args:
            with self.builder.goto_entry_block():
                ptr = self.builder.alloca(arg.type)
            self.builder.store(arg, ptr)
            kernel_args_ptrs.append(ptr)

        # Populate the args_list and the args_ty_list LLVM arrays
        self._populate_kernel_args_and_args_ty_arrays(
            callargs_ptrs=kernel_args_ptrs,
            kernel_argtys=ty_kernel_args,
            args_list=args_list,
            args_ty_list=args_ty_list,
        )

        self.arguments.arg_list = args_list
        self.arguments.arg_ty_list = args_ty_list
        self.arguments.total_kernel_args = self.context.get_constant(
            types.uintp, num_flattened_kernel_args
        )

    def set_dependant_event_list(self):
        """Sets dependant events to the argument list."""
        # TODO: implement for non zero input
        if self.arguments.dep_events is not None:
            return

        self.arguments.dep_events = self.builder.bitcast(
            utils.create_null_ptr(builder=self.builder, context=self.context),
            utils.get_llvm_type(context=self.context, type=types.voidptr),
        )
        self.arguments.dep_events_len = self.context.get_constant(
            types.uintp, 0
        )

    def submit(self) -> llvmir.Instruction:
        """Submits kernel by calling sycl.dpctl_queue_submit_range or
        sycl.dpctl_queue_submit_ndrange. Must be called after all arguments
        set."""
        args = self.arguments.to_list()

        if self.arguments.local_range is None:
            eref = sycl.dpctl_queue_submit_range(self.builder, *args)
        else:
            eref = sycl.dpctl_queue_submit_ndrange(self.builder, *args)

        return eref

    def _get_num_flattened_kernel_args(
        self,
        kernel_argtys: tuple[types.Type, ...],
    ) -> int:
        """Returns number of flattened arguments based on the numba types.
        flattens dpnp arrays and complex values."""
        num_flattened_kernel_args = 0
        for arg_type in kernel_argtys:
            if isinstance(arg_type, DpnpNdArray):
                datamodel = self.context.data_model_manager.lookup(arg_type)
                num_flattened_kernel_args += datamodel.flattened_field_count
            elif arg_type in [types.complex64, types.complex128]:
                num_flattened_kernel_args += 2
            else:
                num_flattened_kernel_args += 1

        return num_flattened_kernel_args

    def _populate_kernel_args_and_args_ty_arrays(
        self,
        kernel_argtys,
        callargs_ptrs,
        args_list,
        args_ty_list,
    ):
        kernel_arg_num = 0
        for arg_num, argtype in enumerate(kernel_argtys):
            llvm_val = callargs_ptrs[arg_num]
            if isinstance(argtype, DpnpNdArray):
                datamodel = self.context.data_model_manager.lookup(argtype)
                self._build_array_arg(
                    array_val=llvm_val,
                    array_data_model=datamodel,
                    array_rank=argtype.ndim,
                    arg_list=args_list,
                    args_ty_list=args_ty_list,
                    arg_num=kernel_arg_num,
                )
                kernel_arg_num += datamodel.flattened_field_count
            else:
                if argtype == types.complex64:
                    self._build_complex_arg(
                        llvm_val,
                        types.float32,
                        args_list,
                        args_ty_list,
                        kernel_arg_num,
                    )
                    kernel_arg_num += 2
                elif argtype == types.complex128:
                    self._build_complex_arg(
                        llvm_val,
                        types.float64,
                        args_list,
                        args_ty_list,
                        kernel_arg_num,
                    )
                    kernel_arg_num += 2
                else:
                    self._build_arg(
                        llvm_val,
                        argtype,
                        args_list,
                        args_ty_list,
                        kernel_arg_num,
                    )
                    kernel_arg_num += 1

1	# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	"""Module that contains numba style wrapper around sycl kernel submit."""	1✔
6
7	from dataclasses import dataclass	1✔
8
9	import dpctl	1✔
10	from llvmlite import ir as llvmir	1✔
11	from llvmlite.ir.builder import IRBuilder	1✔
12	from numba.core import cgutils, types	1✔
13
14	from numba_dpex import config, utils	1✔
15	from numba_dpex.core.runtime.context import DpexRTContext	1✔
16	from numba_dpex.core.types import DpnpNdArray	1✔
17	from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl	1✔
18	from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum	1✔
19
20	MAX_SIZE_OF_SYCL_RANGE = 3	1✔
21
22
23	@dataclass	1✔
24	class _KernelLaunchIRArguments: # pylint: disable=too-many-instance-attributes	1✔
25	"""List of kernel launch arguments used in sycl.dpctl_queue_submit_range and
26	sycl.dpctl_queue_submit_ndrange."""
27
28	sycl_kernel_ref: llvmir.Instruction = None	1✔
29	sycl_queue_ref: llvmir.Instruction = None	1✔
30	arg_list: llvmir.Instruction = None	1✔
31	arg_ty_list: llvmir.Instruction = None	1✔
32	total_kernel_args: llvmir.Instruction = None	1✔
33	global_range: llvmir.Instruction = None	1✔
34	local_range: llvmir.Instruction = None	1✔
35	range_size: llvmir.Instruction = None	1✔
36	dep_events: llvmir.Instruction = None	1✔
37	dep_events_len: llvmir.Instruction = None	1✔
38
39	def to_list(self):	1✔
40	"""Returns list of arguments in the right order to pass to
41	sycl.dpctl_queue_submit_range or sycl.dpctl_queue_submit_ndrange."""
42	res = [	1✔
43	self.sycl_kernel_ref,
44	self.sycl_queue_ref,
45	self.arg_list,
46	self.arg_ty_list,
47	self.total_kernel_args,
48	self.global_range,
49	]
50
51	if self.local_range is not None:	1✔
52	res += [self.local_range]	1✔
53
54	res += [	1✔
55	self.range_size,
56	self.dep_events,
57	self.dep_events_len,
58	]
59
60	return res	1✔
61
62
63	class KernelLaunchIRBuilder:	1✔
64	"""
65	KernelLaunchIRBuilder(lowerer, cres)
66	Helper class to build the LLVM IR for the submission of a kernel.
67
68	The class generates LLVM IR inside the current LLVM module that is needed
69	for submitting kernels. The LLVM Values that
70	"""
71
72	def __init__(self, context, builder: IRBuilder):	1✔
73	"""Create a KernelLauncher for the specified kernel.
74
75	Args:
76	context: A Numba target context that will be used to generate the
77	code.
78	builder: An llvmlite IRBuilder instance used to generate LLVM IR.
79	"""
80	self.context = context	1✔
81	self.builder = builder	1✔
82	self.rtctx = DpexRTContext(self.context)	1✔
83	self.arguments = _KernelLaunchIRArguments()	1✔
84
85	def _build_nullptr(self):	1✔
86	"""Builds the LLVM IR to represent a null pointer.
87
88	Returns: An LLVM Value storing a null pointer
89	"""
90	zero = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)	1✔
91	self.builder.store(self.context.get_constant(types.int64, 0), zero)	1✔
92	return self.builder.bitcast(	1✔
93	zero, utils.get_llvm_type(context=self.context, type=types.voidptr)
94	)
95
96	def _build_array_attr_arg( # pylint: disable=too-many-arguments	1✔
97	self,
98	array_val,
99	array_attr_pos,
100	array_attr_ty,
101	arg_list,
102	args_ty_list,
103	arg_num,
104	):
105	array_attr = self.builder.gep(	1✔
106	array_val,
107	[
108	self.context.get_constant(types.int32, 0),
109	self.context.get_constant(types.int32, array_attr_pos),
110	],
111	)
112
113	# FIXME: If pointer arg then load it to some value and pass that value.
114	# We also most likely need an address space cast
115	if isinstance(	1✔
116	array_attr_ty, (types.misc.RawPointer, types.misc.CPointer)
117	):
118	array_attr = self.builder.load(array_attr)	1✔
119
120	self._build_arg(	1✔
121	val=array_attr,
122	ty=array_attr_ty,
123	arg_list=arg_list,
124	args_ty_list=args_ty_list,
125	arg_num=arg_num,
126	)
127
128	def _build_unituple_member_arg( # pylint: disable=too-many-arguments	1✔
129	self, array_val, array_attr_pos, ndims, arg_list, args_ty_list, arg_num
130	):
131	array_attr = self.builder.gep(	1✔
132	array_val,
133	[
134	self.context.get_constant(types.int32, 0),
135	self.context.get_constant(types.int32, array_attr_pos),
136	],
137	)
138
139	for ndim in range(ndims):	1✔
140	self._build_array_attr_arg(	1✔
141	array_val=array_attr,
142	array_attr_pos=ndim,
143	array_attr_ty=types.int64,
144	arg_list=arg_list,
145	args_ty_list=args_ty_list,
146	arg_num=arg_num + ndim,
147	)
148
149	def _build_arg(	1✔
150	self, val, ty, arg_list, args_ty_list, arg_num
151	): # pylint: disable=too-many-arguments
152	"""Stores the kernel arguments and the kernel argument types into
153	arrays that will be passed to DPCTLQueue_SubmitRange.
154
155	Args:
156	val: An LLVM IR Value that will be stored into the arguments array
157	ty: A Numba type that will be converted to a DPCTLKernelArgType
158	enum and stored into the argument types list array
159	arg_list: An LLVM IR Value array that stores the kernel arguments
160	args_ty_list: An LLVM IR Value array that stores the
161	DPCTLKernelArgType enum value for each kernel argument
162	arg_num: The index position at which the arg_list and args_ty_list
163	need to be updated.
164	"""
165	kernel_arg_dst = self.builder.gep(	1✔
166	arg_list,
167	[self.context.get_constant(types.int32, arg_num)],
168	)
169	kernel_arg_ty_dst = self.builder.gep(	1✔
170	args_ty_list,
171	[self.context.get_constant(types.int32, arg_num)],
172	)
173	val = self.builder.bitcast(	1✔
174	val,
175	utils.get_llvm_type(context=self.context, type=types.voidptr),
176	)
177	self.builder.store(val, kernel_arg_dst)	1✔
178	self.builder.store(	1✔
179	numba_type_to_dpctl_typenum(self.context, ty), kernel_arg_ty_dst
180	)
181
182	def _build_complex_arg(	1✔
183	self, val, ty, arg_list, args_ty_list, arg_num
184	): # pylint: disable=too-many-arguments
185	"""Creates a list of LLVM Values for an unpacked complex kernel
186	argument.
187	"""
188	self._build_array_attr_arg(	1✔
189	array_val=val,
190	array_attr_pos=0,
191	array_attr_ty=ty,
192	arg_list=arg_list,
193	args_ty_list=args_ty_list,
194	arg_num=arg_num,
195	)
196	arg_num += 1	1✔
197	self._build_array_attr_arg(	1✔
198	array_val=val,
199	array_attr_pos=1,
200	array_attr_ty=ty,
201	arg_list=arg_list,
202	args_ty_list=args_ty_list,
203	arg_num=arg_num,
204	)
205	arg_num += 1	1✔
206
207	def _build_array_arg( # pylint: disable=too-many-arguments	1✔
208	self,
209	array_val,
210	array_data_model,
211	array_rank, # pylint: disable=unused-argument
212	arg_list,
213	args_ty_list,
214	arg_num,
215	):
216	"""Creates a list of LLVM Values for an unpacked DpnpNdArray kernel
217	argument.
218
219	The steps performed here are the same as in
220	numba_dpex.core.kernel_interface.arg_pack_unpacker._unpack_array_helper
221	"""
222	# Argument 1: Null pointer for the NRT_MemInfo attribute of the array
223	nullptr = self._build_nullptr()	1✔
224	self._build_arg(	1✔
225	val=nullptr,
226	ty=types.int64,
227	arg_list=arg_list,
228	args_ty_list=args_ty_list,
229	arg_num=arg_num,
230	)
231	arg_num += 1	1✔
232	# Argument 2: Null pointer for the Parent attribute of the array
233	nullptr = self._build_nullptr()	1✔
234	self._build_arg(	1✔
235	val=nullptr,
236	ty=types.int64,
237	arg_list=arg_list,
238	args_ty_list=args_ty_list,
239	arg_num=arg_num,
240	)
241	arg_num += 1	1✔
242	# Argument nitems
243	self._build_array_attr_arg(	1✔
244	array_val=array_val,
245	array_attr_pos=array_data_model.get_field_position("nitems"),
246	array_attr_ty=array_data_model.get_member_fe_type("nitems"),
247	arg_list=arg_list,
248	args_ty_list=args_ty_list,
249	arg_num=arg_num,
250	)
251	arg_num += 1	1✔
252	# Argument itemsize
253	self._build_array_attr_arg(	1✔
254	array_val=array_val,
255	array_attr_pos=array_data_model.get_field_position("itemsize"),
256	array_attr_ty=array_data_model.get_member_fe_type("itemsize"),
257	arg_list=arg_list,
258	args_ty_list=args_ty_list,
259	arg_num=arg_num,
260	)
261	arg_num += 1	1✔
262	# Argument data
263	self._build_array_attr_arg(	1✔
264	array_val=array_val,
265	array_attr_pos=array_data_model.get_field_position("data"),
266	array_attr_ty=array_data_model.get_member_fe_type("data"),
267	arg_list=arg_list,
268	args_ty_list=args_ty_list,
269	arg_num=arg_num,
270	)
271	arg_num += 1	1✔
272	# Argument sycl_queue: as the queue pointer is not to be used in a
273	# kernel we always pass in a nullptr
274	self._build_arg(	1✔
275	val=nullptr,
276	ty=types.int64,
277	arg_list=arg_list,
278	args_ty_list=args_ty_list,
279	arg_num=arg_num,
280	)
281	arg_num += 1	1✔
282	# Arguments for shape
283	shape_member = array_data_model.get_member_fe_type("shape")	1✔
284	self._build_unituple_member_arg(	1✔
285	array_val=array_val,
286	array_attr_pos=array_data_model.get_field_position("shape"),
287	ndims=shape_member.count,
288	arg_list=arg_list,
289	args_ty_list=args_ty_list,
290	arg_num=arg_num,
291	)
292	arg_num += shape_member.count	1✔
293	# Arguments for strides
294	stride_member = array_data_model.get_member_fe_type("strides")	1✔
295	self._build_unituple_member_arg(	1✔
296	array_val=array_val,
297	array_attr_pos=array_data_model.get_field_position("strides"),
298	ndims=stride_member.count,
299	arg_list=arg_list,
300	args_ty_list=args_ty_list,
301	arg_num=arg_num,
302	)
303	arg_num += stride_member.count	1✔
304
305	# TODO: remove, not part of the builder
306	def get_queue(self, exec_queue: dpctl.SyclQueue) -> llvmir.Instruction:	1✔
307	"""Allocates memory on the stack to store a DPCTLSyclQueueRef.
308
309	Returns: A LLVM Value storing the pointer to the SYCL queue created
310	using the filter string for the Python exec_queue (dpctl.SyclQueue).
311	"""
312
313	# Allocate a stack var to store the queue created from the filter string
314	sycl_queue_val = cgutils.alloca_once(	1✔
315	self.builder,
316	utils.get_llvm_type(context=self.context, type=types.voidptr),
317	)
318	# Insert a global constant to store the filter string
319	device = self.context.insert_const_string(	1✔
320	self.builder.module, exec_queue.sycl_device.filter_string
321	)
322	# Store the queue returned by DPEXRTQueue_CreateFromFilterString in a
323	# local variable
324	self.builder.store(	1✔
325	self.rtctx.get_queue_from_filter_string(
326	builder=self.builder, device=device
327	),
328	sycl_queue_val,
329	)
330	return self.builder.load(sycl_queue_val)	1✔
331
332	def _allocate_kernel_arg_array(self, num_kernel_args):	1✔
333	"""Allocates an array to store the LLVM Value for every kernel argument.
334
335	Args:
336	num_kernel_args (int): The number of kernel arguments that
337	determines the size of args array to allocate.
338
339	Returns: An LLVM IR value pointing to an array to store the kernel
340	arguments.
341	"""
342	args_list = cgutils.alloca_once(	1✔
343	self.builder,
344	utils.LLVMTypes.byte_ptr_t,
345	size=self.context.get_constant(types.uintp, num_kernel_args),
346	)
347
348	return args_list	1✔
349
350	def _allocate_kernel_arg_ty_array(self, num_kernel_args):	1✔
351	"""Allocates an array to store the LLVM Value for the typenum for
352	every kernel argument.
353
354	Args:
355	num_kernel_args (int): The number of kernel arguments that
356	determines the size of args array to allocate.
357
358	Returns: An LLVM IR value pointing to an array to store the kernel
359	arguments typenums as defined in dpctl.
360	"""
361	args_ty_list = cgutils.alloca_once(	1✔
362	self.builder,
363	utils.LLVMTypes.int32_t,
364	size=self.context.get_constant(types.uintp, num_kernel_args),
365	)
366
367	return args_ty_list	1✔
368
369	def _create_sycl_range(self, idx_range):	1✔
370	"""Allocate a size_t[3] array to store the extents of a sycl::range.
371
372	Sycl supports upto 3-dimensional ranges and a such the array is
373	statically sized to length three. Only the elements that store an actual
374	range value are populated based on the size of the idx_range argument.
375
376	"""
377	intp_t = utils.get_llvm_type(context=self.context, type=types.intp)	1✔
378	intp_ptr_t = utils.get_llvm_ptr_type(intp_t)	1✔
379	num_dim = len(idx_range)	1✔
380
381	# form the global range
382	range_list = cgutils.alloca_once(	1✔
383	self.builder,
384	utils.get_llvm_type(context=self.context, type=types.uintp),
385	size=self.context.get_constant(types.uintp, MAX_SIZE_OF_SYCL_RANGE),
386	)
387
388	for i in range(num_dim):	1✔
389	rext = idx_range[i]	1✔
390	if rext.type != utils.LLVMTypes.int64_t:	1!
UNCOV 391	rext = self.builder.sext(rext, utils.LLVMTypes.int64_t)	×
392
393	# we reverse the global range to account for how sycl and opencl
394	# range differs
395	self.builder.store(	1✔
396	rext,
397	self.builder.gep(
398	range_list,
399	[self.context.get_constant(types.uintp, (num_dim - 1) - i)],
400	),
401	)
402
403	return self.builder.bitcast(range_list, intp_ptr_t)	1✔
404
405	def set_kernel(self, sycl_kernel_ref: llvmir.Instruction):	1✔
406	"""Sets kernel to the argument list."""
407	self.arguments.sycl_kernel_ref = sycl_kernel_ref	1✔
408
409	def set_queue(self, sycl_queue_ref: llvmir.Instruction):	1✔
410	"""Sets queue to the argument list."""
411	self.arguments.sycl_queue_ref = sycl_queue_ref	1✔
412
413	def set_range(	1✔
414	self,
415	global_range: list,
416	local_range: list = None,
417	):
418	"""Sets global and local range if provided to the argument list."""
419	self.arguments.global_range = self._create_sycl_range(global_range)	1✔
420	if local_range is not None and len(local_range) > 0:	1✔
421	self.arguments.local_range = self._create_sycl_range(local_range)	1✔
422	self.arguments.range_size = self.context.get_constant(	1✔
423	types.uintp, len(global_range)
424	)
425
426	def set_arguments(	1✔
427	self,
428	ty_kernel_args: list,
429	kernel_args: list,
430	):
431	"""Sets flattened kernel args, kernel arg types and number of those
432	arguments to the argument list."""
433	if config.DEBUG_KERNEL_LAUNCHER:	1!
UNCOV 434	cgutils.printf(	×
435	self.builder,
436	"DPEX-DEBUG: Populating kernel args and arg type arrays.\n",
437	)
438
439	num_flattened_kernel_args = self._get_num_flattened_kernel_args(	1✔
440	kernel_argtys=ty_kernel_args,
441	)
442
443	# Create LLVM values for the kernel args list and kernel arg types list
444	args_list = self._allocate_kernel_arg_array(num_flattened_kernel_args)	1✔
445
446	args_ty_list = self._allocate_kernel_arg_ty_array(	1✔
447	num_flattened_kernel_args
448	)
449
450	kernel_args_ptrs = []	1✔
451	for arg in kernel_args:	1✔
452	with self.builder.goto_entry_block():	1✔
453	ptr = self.builder.alloca(arg.type)	1✔
454	self.builder.store(arg, ptr)	1✔
455	kernel_args_ptrs.append(ptr)	1✔
456
457	# Populate the args_list and the args_ty_list LLVM arrays
458	self._populate_kernel_args_and_args_ty_arrays(	1✔
459	callargs_ptrs=kernel_args_ptrs,
460	kernel_argtys=ty_kernel_args,
461	args_list=args_list,
462	args_ty_list=args_ty_list,
463	)
464
465	self.arguments.arg_list = args_list	1✔
466	self.arguments.arg_ty_list = args_ty_list	1✔
467	self.arguments.total_kernel_args = self.context.get_constant(	1✔
468	types.uintp, num_flattened_kernel_args
469	)
470
471	def set_dependant_event_list(self):	1✔
472	"""Sets dependant events to the argument list."""
473	# TODO: implement for non zero input
474	if self.arguments.dep_events is not None:	1!
NEW 475	return	×
476
477	self.arguments.dep_events = self.builder.bitcast(	1✔
478	utils.create_null_ptr(builder=self.builder, context=self.context),
479	utils.get_llvm_type(context=self.context, type=types.voidptr),
480	)
481	self.arguments.dep_events_len = self.context.get_constant(	1✔
482	types.uintp, 0
483	)
484
485	def submit(self) -> llvmir.Instruction:	1✔
486	"""Submits kernel by calling sycl.dpctl_queue_submit_range or
487	sycl.dpctl_queue_submit_ndrange. Must be called after all arguments
488	set."""
489	args = self.arguments.to_list()	1✔
490
491	if self.arguments.local_range is None:	1✔
492	eref = sycl.dpctl_queue_submit_range(self.builder, *args)	1✔
493	else:
494	eref = sycl.dpctl_queue_submit_ndrange(self.builder, *args)	1✔
495
496	return eref	1✔
497
498	def _get_num_flattened_kernel_args(	1✔
499	self,
500	kernel_argtys: tuple[types.Type, ...],
501	) -> int:
502	"""Returns number of flattened arguments based on the numba types.
503	flattens dpnp arrays and complex values."""
504	num_flattened_kernel_args = 0	1✔
505	for arg_type in kernel_argtys:	1✔
506	if isinstance(arg_type, DpnpNdArray):	1✔
507	datamodel = self.context.data_model_manager.lookup(arg_type)	1✔
508	num_flattened_kernel_args += datamodel.flattened_field_count	1✔
509	elif arg_type in [types.complex64, types.complex128]:	1✔
510	num_flattened_kernel_args += 2	1✔
511	else:
512	num_flattened_kernel_args += 1	1✔
513
514	return num_flattened_kernel_args	1✔
515
516	def _populate_kernel_args_and_args_ty_arrays(	1✔
517	self,
518	kernel_argtys,
519	callargs_ptrs,
520	args_list,
521	args_ty_list,
522	):
523	kernel_arg_num = 0	1✔
524	for arg_num, argtype in enumerate(kernel_argtys):	1✔
525	llvm_val = callargs_ptrs[arg_num]	1✔
526	if isinstance(argtype, DpnpNdArray):	1✔
527	datamodel = self.context.data_model_manager.lookup(argtype)	1✔
528	self._build_array_arg(	1✔
529	array_val=llvm_val,
530	array_data_model=datamodel,
531	array_rank=argtype.ndim,
532	arg_list=args_list,
533	args_ty_list=args_ty_list,
534	arg_num=kernel_arg_num,
535	)
536	kernel_arg_num += datamodel.flattened_field_count	1✔
537	else:
538	if argtype == types.complex64:	1✔
539	self._build_complex_arg(	1✔
540	llvm_val,
541	types.float32,
542	args_list,
543	args_ty_list,
544	kernel_arg_num,
545	)
546	kernel_arg_num += 2	1✔
547	elif argtype == types.complex128:	1✔
548	self._build_complex_arg(	1✔
549	llvm_val,
550	types.float64,
551	args_list,
552	args_ty_list,
553	kernel_arg_num,
554	)
555	kernel_arg_num += 2	1✔
556	else:
557	self._build_arg(	1✔
558	llvm_val,
559	argtype,
560	args_list,
561	args_ty_list,
562	kernel_arg_num,
563	)
564	kernel_arg_num += 1	1✔

IntelPython / numba-dpex / 7201235944

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