7187743552

Committed 12 Dec 2023 09:58PM UTC coverage: 79.164% (-0.08%) from 79.239%

Build # 7187743552

Build Type

Pull #1144

github

Committed by

web-flow

Commit Message

Merge 3c710c7af into 8fc7fbddb

Pull Request Pull Request #1144: Implement support for dpnp.arange()

Run Details

1261 of 1863 branches covered (0.0%)

Branch coverage included in aggregate %.

148 of 183 new or added lines in 3 files covered. (80.87%)

1 existing line in 1 file now uncovered.

5464 of 6632 relevant lines covered (82.39%)

0.82 hits per line

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

87.78

/numba_dpex/core/parfors/parfor_lowerer.py

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

import copy
from collections import namedtuple

from llvmlite import ir as llvmir
from numba.core import cgutils, ir, types
from numba.parfors.parfor import (
    find_potential_aliases_parfor,
    get_parfor_outputs,
)

from numba_dpex import config
from numba_dpex.core.datamodel.models import dpex_data_model_manager as dpex_dmm
from numba_dpex.core.parfors.reduction_helper import (
    ReductionHelper,
    ReductionKernelVariables,
)
from numba_dpex.core.utils.kernel_launcher import KernelLaunchIRBuilder

from ..exceptions import UnsupportedParforError
from ..types.dpnp_ndarray_type import DpnpNdArray
from .kernel_builder import create_kernel_for_parfor
from .reduction_kernel_builder import (
    create_reduction_main_kernel_for_parfor,
    create_reduction_remainder_kernel_for_parfor,
)

_KernelArgs = namedtuple(
    "_KernelArgs",
    ["num_flattened_args", "arg_vals", "arg_types"],
)


# A global list of kernels to keep the objects alive indefinitely.
keep_alive_kernels = []


def _getvar(lowerer, x):
    """Returns the LLVM Value corresponding to a Numba IR variable.

    Depending on if Numba's sroa-like optimization is enabled or not, the
    LLVM value for an Numba IR variable is found in either the ``varmap``
    or the ``blk_local_varmap`` of the ``lowerer``. If the LLVM Value is not a
    pointer, e.g., in case of function args with sroa optimization enabled, then
    creates an alloca and stores the Value into the new alloca Value and returns
    it. The extra alloca is needed as all inputs to a kernel function need to
    be passed by reference and not value.

    Args:
        lowerer: The Numba Lower instance used to lower the function.
        x: Numba IR variable name used to lookup the corresponding
           LLVM Value.

    Raises:
        AssertionError: If the LLVM Value for ``x`` does not exist in either
            the ``varmap`` or the ``blk_local_varmap``.

    Returns: An LLVM Value object

    """
    var_val = None
    if x in lowerer._blk_local_varmap:
        var_val = lowerer._blk_local_varmap[x]
    elif x in lowerer.varmap:
        var_val = lowerer.varmap[x]

    if var_val:
        if not isinstance(var_val.type, llvmir.PointerType):
            with lowerer.builder.goto_entry_block():
                var_val_ptr = lowerer.builder.alloca(var_val.type)
            lowerer.builder.store(var_val, var_val_ptr)
            return var_val_ptr
        else:
            return var_val
    else:
        raise AssertionError("No llvm Value found for kernel arg")


def _load_range(lowerer, value):
    if isinstance(value, ir.Var):
        return lowerer.loadvar(value.name)
    else:
        return lowerer.context.get_constant(types.uintp, value)


class ParforLowerImpl:
    """Provides a custom lowerer for parfor nodes that generates a SYCL kernel
    for a parfor and submits it to a queue.
    """

    def _build_kernel_arglist(self, kernel_fn, lowerer, kernel_builder):
        """Creates local variables for all the arguments and the argument types
        that are passes to the kernel function.

        Args:
            kernel_fn: Kernel function to be launched.
            lowerer: The Numba lowerer used to generate the LLVM IR

        Raises:
            AssertionError: If the LLVM IR Value for an argument defined in
            Numba IR is not found.
        """
        num_flattened_args = 0

        # Compute number of args to be passed to the kernel. Note that the
        # actual number of kernel arguments is greater than the count of
        # kernel_fn.kernel_args as arrays get flattened.
        for arg_type in kernel_fn.kernel_arg_types:
            if isinstance(arg_type, DpnpNdArray):
                datamodel = dpex_dmm.lookup(arg_type)
                num_flattened_args += datamodel.flattened_field_count
            elif arg_type == types.complex64 or arg_type == types.complex128:
                num_flattened_args += 2
            else:
                num_flattened_args += 1

        # Create LLVM values for the kernel args list and kernel arg types list
        args_list = kernel_builder.allocate_kernel_arg_array(num_flattened_args)
        args_ty_list = kernel_builder.allocate_kernel_arg_ty_array(
            num_flattened_args
        )
        callargs_ptrs = []
        for arg in kernel_fn.kernel_args:
            callargs_ptrs.append(_getvar(lowerer, arg))

        kernel_builder.populate_kernel_args_and_args_ty_arrays(
            kernel_argtys=kernel_fn.kernel_arg_types,
            callargs_ptrs=callargs_ptrs,
            args_list=args_list,
            args_ty_list=args_ty_list,
        )

        return _KernelArgs(
            num_flattened_args=num_flattened_args,
            arg_vals=args_list,
            arg_types=args_ty_list,
        )

    def _submit_parfor_kernel(
        self,
        lowerer,
        kernel_fn,
        loop_ranges,
    ):
        """
        Adds a call to submit a kernel function into the function body of the
        current Numba JIT compiled function.
        """
        # Ensure that the Python arguments are kept alive for the duration of
        # the kernel execution
        keep_alive_kernels.append(kernel_fn.kernel)
        kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)

        ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)
        args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)

        # Create a global range over which to submit the kernel based on the
        # loop_ranges of the parfor
        global_range = []
        # SYCL ranges can have at max 3 dimension. If the parfor is of a higher
        # dimension then the indexing for the higher dimensions is done inside
        # the kernel.
        global_range_rank = len(loop_ranges) if len(loop_ranges) < 3 else 3

        for i in range(global_range_rank):
            start, stop, step = loop_ranges[i]
            stop = _load_range(lowerer, stop)
            if step != 1:
                raise UnsupportedParforError(
                    "non-unit strides are not yet supported."
                )
            global_range.append(stop)

        local_range = []

        kernel_ref_addr = kernel_fn.kernel.addressof_ref()
        kernel_ref = lowerer.builder.inttoptr(
            lowerer.context.get_constant(types.uintp, kernel_ref_addr),
            cgutils.voidptr_t,
        )
        curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)

        # Submit a synchronous kernel
        kernel_builder.submit_sycl_kernel(
            sycl_kernel_ref=kernel_ref,
            sycl_queue_ref=curr_queue_ref,
            total_kernel_args=args.num_flattened_args,
            arg_list=args.arg_vals,
            arg_ty_list=args.arg_types,
            global_range=global_range,
            local_range=local_range,
        )

        # At this point we can free the DPCTLSyclQueueRef (curr_queue)
        kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)

    def _submit_reduction_main_parfor_kernel(
        self,
        lowerer,
        kernel_fn,
        reductionHelper=None,
    ):
        """
        Adds a call to submit the main kernel of a parfor reduction into the
        function body of the current Numba JIT compiled function.
        """
        # Ensure that the Python arguments are kept alive for the duration of
        # the kernel execution
        keep_alive_kernels.append(kernel_fn.kernel)
        kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)

        ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)

        args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)
        # Create a global range over which to submit the kernel based on the
        # loop_ranges of the parfor
        global_range = []

        stop = reductionHelper.global_size_var
        stop = _load_range(lowerer, stop)
        global_range.append(stop)

        local_range = []
        local_range.append(
            _load_range(lowerer, reductionHelper.work_group_size)
        )

        kernel_ref_addr = kernel_fn.kernel.addressof_ref()
        kernel_ref = lowerer.builder.inttoptr(
            lowerer.context.get_constant(types.uintp, kernel_ref_addr),
            cgutils.voidptr_t,
        )
        curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)

        # Submit a synchronous kernel
        kernel_builder.submit_sycl_kernel(
            sycl_kernel_ref=kernel_ref,
            sycl_queue_ref=curr_queue_ref,
            total_kernel_args=args.num_flattened_args,
            arg_list=args.arg_vals,
            arg_ty_list=args.arg_types,
            global_range=global_range,
            local_range=local_range,
        )

        # At this point we can free the DPCTLSyclQueueRef (curr_queue)
        kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)

    def _submit_reduction_remainder_parfor_kernel(
        self,
        lowerer,
        kernel_fn,
    ):
        """
        Adds a call to submit the remainder kernel of a parfor reduction into
        the function body of the current Numba JIT compiled function.
        """
        # Ensure that the Python arguments are kept alive for the duration of
        # the kernel execution
        keep_alive_kernels.append(kernel_fn.kernel)

        kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)

        ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)

        args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)
        # Create a global range over which to submit the kernel based on the
        # loop_ranges of the parfor
        global_range = []

        stop = _load_range(lowerer, 1)

        global_range.append(stop)

        local_range = []

        kernel_ref_addr = kernel_fn.kernel.addressof_ref()
        kernel_ref = lowerer.builder.inttoptr(
            lowerer.context.get_constant(types.uintp, kernel_ref_addr),
            cgutils.voidptr_t,
        )
        curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)

        # Submit a synchronous kernel
        kernel_builder.submit_sycl_kernel(
            sycl_kernel_ref=kernel_ref,
            sycl_queue_ref=curr_queue_ref,
            total_kernel_args=args.num_flattened_args,
            arg_list=args.arg_vals,
            arg_ty_list=args.arg_types,
            global_range=global_range,
            local_range=local_range,
        )

        # At this point we can free the DPCTLSyclQueueRef (curr_queue)
        kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)

    def _reduction_codegen(
        self,
        parfor,
        typemap,
        nredvars,
        parfor_redvars,
        parfor_reddict,
        lowerer,
        parfor_output_arrays,
        loop_ranges,
        flags,
        alias_map,
    ):
        """
        Reduction kernel generation and submission.
        """
        inputArrayName = None
        for para in parfor.params:
            if not isinstance(typemap[para], DpnpNdArray):
                continue
            inputArrayName = para
            break

        if inputArrayName is None:
            raise AssertionError

        inputArrayType = typemap[inputArrayName]

        reductionHelperList = []
        for i in range(nredvars):
            reductionHelper = ReductionHelper()
            reductionHelper._allocate_partial_reduction_arrays(
                parfor,
                lowerer,
                parfor_redvars[i],
                inputArrayType,
            )
            reductionHelperList.append(reductionHelper)

        reductionKernelVar = ReductionKernelVariables(
            lowerer=lowerer,
            parfor_node=parfor,
            typemap=typemap,
            parfor_outputs=parfor_output_arrays,
            reductionHelperList=reductionHelperList,
        )

        if len(parfor.loop_nests) > 1:
            raise UnsupportedParforError(
                "Reduction with nested loop is not yet supported."
            )

        parfor_kernel = create_reduction_main_kernel_for_parfor(
            loop_ranges,
            parfor,
            typemap,
            flags,
            bool(alias_map),
            reductionKernelVar,
            parfor_reddict,
        )

        self._submit_reduction_main_parfor_kernel(
            lowerer,
            parfor_kernel,
            reductionHelperList[0],
        )

        parfor_kernel = create_reduction_remainder_kernel_for_parfor(
            parfor,
            typemap,
            flags,
            bool(alias_map),
            reductionKernelVar,
            parfor_reddict,
            reductionHelperList,
        )

        self._submit_reduction_remainder_parfor_kernel(
            lowerer,
            parfor_kernel,
        )

        reductionKernelVar.copy_final_sum_to_host(parfor_kernel)

    def _lower_parfor_as_kernel(self, lowerer, parfor):
        """Lowers a parfor node created by the dpjit compiler to a
        ``numba_dpex.kernel``.

        The general approach is as follows:

            - The code from the parfor's init block is lowered normally
              in the context of the current function.
            - The body of the parfor is transformed into a kernel function.
            - Dpctl runtime calls to submit the kernel are added.

        """
        # We copy the typemap here because for race condition variable we'll
        # update their type to array so they can be updated by the kernel.
        orig_typemap = lowerer.fndesc.typemap

        # replace original typemap with copy and restore the original at the
        # end.
        lowerer.fndesc.typemap = copy.copy(orig_typemap)

        if config.DEBUG_ARRAY_OPT:
            print("lowerer.fndesc", lowerer.fndesc, type(lowerer.fndesc))

        typemap = lowerer.fndesc.typemap
        varmap = lowerer.varmap

        loc = parfor.init_block.loc
        scope = parfor.init_block.scope

        # Lower the init block of the parfor.
        for instr in parfor.init_block.body:
            lowerer.lower_inst(instr)

        for racevar in parfor.races:
            if racevar not in varmap:
                rvtyp = typemap[racevar]
                rv = ir.Var(scope, racevar, loc)
                lowerer._alloca_var(rv.name, rvtyp)

        alias_map = {}
        arg_aliases = {}

        find_potential_aliases_parfor(
            parfor,
            parfor.params,
            typemap,
            lowerer.func_ir,
            alias_map,
            arg_aliases,
        )

        # run get_parfor_outputs() and get_parfor_reductions() before
        # kernel creation since Jumps are modified so CFG of loop_body
        # dict will become invalid
        if parfor.params is None:
            raise AssertionError

        parfor_output_arrays = get_parfor_outputs(parfor, parfor.params)

        # compile parfor body as a separate dpex kernel function
        flags = copy.copy(parfor.flags)
        flags.error_model = "numpy"

        # Can't get here unless
        # flags.set('auto_parallel', ParallelOptions(True)) # noqa: E800
        index_var_typ = typemap[parfor.loop_nests[0].index_variable.name]

        # index variables should have the same type, check rest of indices
        for loop_nest in parfor.loop_nests[1:]:
            if typemap[loop_nest.index_variable.name] != index_var_typ:
                raise AssertionError

        loop_ranges = [
            (loop_nest.start, loop_nest.stop, loop_nest.step)
            for loop_nest in parfor.loop_nests
        ]

        parfor_redvars, parfor_reddict = parfor.redvars, parfor.reddict

        nredvars = len(parfor_redvars)
        if nredvars > 0:
            self._reduction_codegen(
                parfor,
                typemap,
                nredvars,
                parfor_redvars,
                parfor_reddict,
                lowerer,
                parfor_output_arrays,
                loop_ranges,
                flags,
                alias_map,
            )
        else:
            try:
                parfor_kernel = create_kernel_for_parfor(
                    lowerer,
                    parfor,
                    typemap,
                    flags,
                    loop_ranges,
                    bool(alias_map),
                    parfor.races,
                    parfor_output_arrays,
                )
            except Exception:
                # FIXME: Make the exception more informative
                raise UnsupportedParforError

            # Finally submit the kernel
            self._submit_parfor_kernel(
                lowerer,
                parfor_kernel,
                loop_ranges,
            )

        # TODO: free the kernel at this point

        # Restore the original typemap of the function that was replaced
        # temporarily at the beginning of this function.
        lowerer.fndesc.typemap = orig_typemap


class ParforLowerFactory:
    """A pseudo-factory class that maps a device filter string to a lowering
    function.

    Each Parfor instruction can have an optional "lowerer" attribute. The
    lowerer attribute determines how the parfor instruction should be lowered
    to LLVM IR. In addition, the lower attribute decides which parfor
    instructions can be fused together.

    The factory class maintains a dictionary mapping every device
    type (filter string) encountered so far to a lowerer function for that
    device type. At this point numba-dpex does not generate device-specific code
    and the lowerer used is same for all device types. However, as a different
    ParforLowerImpl instance is returned for every parfor instruction that has
    a distinct compute-follows-data inferred device it prevents illegal
    parfor fusion.
    """

    device_to_lowerer_map = {}

    @classmethod
    def get_lowerer(cls, device):
        try:
            lowerer = ParforLowerFactory.device_to_lowerer_map[device]
        except KeyError:
            lowerer = ParforLowerImpl()._lower_parfor_as_kernel
            ParforLowerFactory.device_to_lowerer_map[device] = lowerer

        return lowerer

1	# SPDX-FileCopyrightText: 2020 - 2022 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import copy	1✔
6	from collections import namedtuple	1✔
7
8	from llvmlite import ir as llvmir	1✔
9	from numba.core import cgutils, ir, types	1✔
10	from numba.parfors.parfor import (	1✔
11	find_potential_aliases_parfor,
12	get_parfor_outputs,
13	)
14
15	from numba_dpex import config	1✔
16	from numba_dpex.core.datamodel.models import dpex_data_model_manager as dpex_dmm	1✔
17	from numba_dpex.core.parfors.reduction_helper import (	1✔
18	ReductionHelper,
19	ReductionKernelVariables,
20	)
21	from numba_dpex.core.utils.kernel_launcher import KernelLaunchIRBuilder	1✔
22
23	from ..exceptions import UnsupportedParforError	1✔
24	from ..types.dpnp_ndarray_type import DpnpNdArray	1✔
25	from .kernel_builder import create_kernel_for_parfor	1✔
26	from .reduction_kernel_builder import (	1✔
27	create_reduction_main_kernel_for_parfor,
28	create_reduction_remainder_kernel_for_parfor,
29	)
30
31	_KernelArgs = namedtuple(	1✔
32	"_KernelArgs",
33	["num_flattened_args", "arg_vals", "arg_types"],
34	)
35
36
37	# A global list of kernels to keep the objects alive indefinitely.
38	keep_alive_kernels = []	1✔
39
40
41	def _getvar(lowerer, x):	1✔
42	"""Returns the LLVM Value corresponding to a Numba IR variable.
43
44	Depending on if Numba's sroa-like optimization is enabled or not, the
45	LLVM value for an Numba IR variable is found in either the ``varmap``
46	or the ``blk_local_varmap`` of the ``lowerer``. If the LLVM Value is not a
47	pointer, e.g., in case of function args with sroa optimization enabled, then
48	creates an alloca and stores the Value into the new alloca Value and returns
49	it. The extra alloca is needed as all inputs to a kernel function need to
50	be passed by reference and not value.
51
52	Args:
53	lowerer: The Numba Lower instance used to lower the function.
54	x: Numba IR variable name used to lookup the corresponding
55	LLVM Value.
56
57	Raises:
58	AssertionError: If the LLVM Value for ``x`` does not exist in either
59	the ``varmap`` or the ``blk_local_varmap``.
60
61	Returns: An LLVM Value object
62
63	"""
64	var_val = None	1✔
65	if x in lowerer._blk_local_varmap:	1✔
66	var_val = lowerer._blk_local_varmap[x]	1✔
67	elif x in lowerer.varmap:	1!
68	var_val = lowerer.varmap[x]	1✔
69
70	if var_val:	1!
71	if not isinstance(var_val.type, llvmir.PointerType):	1✔
72	with lowerer.builder.goto_entry_block():	1✔
73	var_val_ptr = lowerer.builder.alloca(var_val.type)	1✔
74	lowerer.builder.store(var_val, var_val_ptr)	1✔
75	return var_val_ptr	1✔
76	else:
77	return var_val	1✔
78	else:
79	raise AssertionError("No llvm Value found for kernel arg")	×
80
81
82	def _load_range(lowerer, value):	1✔
83	if isinstance(value, ir.Var):	1✔
84	return lowerer.loadvar(value.name)	1✔
85	else:
86	return lowerer.context.get_constant(types.uintp, value)	1✔
87
88
89	class ParforLowerImpl:	1✔
90	"""Provides a custom lowerer for parfor nodes that generates a SYCL kernel
91	for a parfor and submits it to a queue.
92	"""
93
94	def _build_kernel_arglist(self, kernel_fn, lowerer, kernel_builder):	1✔
95	"""Creates local variables for all the arguments and the argument types
96	that are passes to the kernel function.
97
98	Args:
99	kernel_fn: Kernel function to be launched.
100	lowerer: The Numba lowerer used to generate the LLVM IR
101
102	Raises:
103	AssertionError: If the LLVM IR Value for an argument defined in
104	Numba IR is not found.
105	"""
106	num_flattened_args = 0	1✔
107
108	# Compute number of args to be passed to the kernel. Note that the
109	# actual number of kernel arguments is greater than the count of
110	# kernel_fn.kernel_args as arrays get flattened.
111	for arg_type in kernel_fn.kernel_arg_types:	1✔
112	if isinstance(arg_type, DpnpNdArray):	1✔
113	datamodel = dpex_dmm.lookup(arg_type)	1✔
114	num_flattened_args += datamodel.flattened_field_count	1✔
115	elif arg_type == types.complex64 or arg_type == types.complex128:	1✔
116	num_flattened_args += 2	1✔
117	else:
118	num_flattened_args += 1	1✔
119
120	# Create LLVM values for the kernel args list and kernel arg types list
121	args_list = kernel_builder.allocate_kernel_arg_array(num_flattened_args)	1✔
122	args_ty_list = kernel_builder.allocate_kernel_arg_ty_array(	1✔
123	num_flattened_args
124	)
125	callargs_ptrs = []	1✔
126	for arg in kernel_fn.kernel_args:	1✔
127	callargs_ptrs.append(_getvar(lowerer, arg))	1✔
128
129	kernel_builder.populate_kernel_args_and_args_ty_arrays(	1✔
130	kernel_argtys=kernel_fn.kernel_arg_types,
131	callargs_ptrs=callargs_ptrs,
132	args_list=args_list,
133	args_ty_list=args_ty_list,
134	)
135
136	return _KernelArgs(	1✔
137	num_flattened_args=num_flattened_args,
138	arg_vals=args_list,
139	arg_types=args_ty_list,
140	)
141
142	def _submit_parfor_kernel(	1✔
143	self,
144	lowerer,
145	kernel_fn,
146	loop_ranges,
147	):
148	"""
149	Adds a call to submit a kernel function into the function body of the
150	current Numba JIT compiled function.
151	"""
152	# Ensure that the Python arguments are kept alive for the duration of
153	# the kernel execution
154	keep_alive_kernels.append(kernel_fn.kernel)	1✔
155	kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)	1✔
156
157	ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)	1✔
158	args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)	1✔
159
160	# Create a global range over which to submit the kernel based on the
161	# loop_ranges of the parfor
162	global_range = []	1✔
163	# SYCL ranges can have at max 3 dimension. If the parfor is of a higher
164	# dimension then the indexing for the higher dimensions is done inside
165	# the kernel.
166	global_range_rank = len(loop_ranges) if len(loop_ranges) < 3 else 3	1✔
167
168	for i in range(global_range_rank):	1✔
169	start, stop, step = loop_ranges[i]	1✔
170	stop = _load_range(lowerer, stop)	1✔
171	if step != 1:	1!
172	raise UnsupportedParforError(	×
173	"non-unit strides are not yet supported."
174	)
175	global_range.append(stop)	1✔
176
177	local_range = []	1✔
178
179	kernel_ref_addr = kernel_fn.kernel.addressof_ref()	1✔
180	kernel_ref = lowerer.builder.inttoptr(	1✔
181	lowerer.context.get_constant(types.uintp, kernel_ref_addr),
182	cgutils.voidptr_t,
183	)
184	curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)	1✔
185
186	# Submit a synchronous kernel
187	kernel_builder.submit_sycl_kernel(	1✔
188	sycl_kernel_ref=kernel_ref,
189	sycl_queue_ref=curr_queue_ref,
190	total_kernel_args=args.num_flattened_args,
191	arg_list=args.arg_vals,
192	arg_ty_list=args.arg_types,
193	global_range=global_range,
194	local_range=local_range,
195	)
196
197	# At this point we can free the DPCTLSyclQueueRef (curr_queue)
198	kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)	1✔
199
200	def _submit_reduction_main_parfor_kernel(	1✔
201	self,
202	lowerer,
203	kernel_fn,
204	reductionHelper=None,
205	):
206	"""
207	Adds a call to submit the main kernel of a parfor reduction into the
208	function body of the current Numba JIT compiled function.
209	"""
210	# Ensure that the Python arguments are kept alive for the duration of
211	# the kernel execution
212	keep_alive_kernels.append(kernel_fn.kernel)	1✔
213	kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)	1✔
214
215	ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)	1✔
216
217	args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)	1✔
218	# Create a global range over which to submit the kernel based on the
219	# loop_ranges of the parfor
220	global_range = []	1✔
221
222	stop = reductionHelper.global_size_var	1✔
223	stop = _load_range(lowerer, stop)	1✔
224	global_range.append(stop)	1✔
225
226	local_range = []	1✔
227	local_range.append(	1✔
228	_load_range(lowerer, reductionHelper.work_group_size)
229	)
230
231	kernel_ref_addr = kernel_fn.kernel.addressof_ref()	1✔
232	kernel_ref = lowerer.builder.inttoptr(	1✔
233	lowerer.context.get_constant(types.uintp, kernel_ref_addr),
234	cgutils.voidptr_t,
235	)
236	curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)	1✔
237
238	# Submit a synchronous kernel
239	kernel_builder.submit_sycl_kernel(	1✔
240	sycl_kernel_ref=kernel_ref,
241	sycl_queue_ref=curr_queue_ref,
242	total_kernel_args=args.num_flattened_args,
243	arg_list=args.arg_vals,
244	arg_ty_list=args.arg_types,
245	global_range=global_range,
246	local_range=local_range,
247	)
248
249	# At this point we can free the DPCTLSyclQueueRef (curr_queue)
250	kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)	1✔
251
252	def _submit_reduction_remainder_parfor_kernel(	1✔
253	self,
254	lowerer,
255	kernel_fn,
256	):
257	"""
258	Adds a call to submit the remainder kernel of a parfor reduction into
259	the function body of the current Numba JIT compiled function.
260	"""
261	# Ensure that the Python arguments are kept alive for the duration of
262	# the kernel execution
263	keep_alive_kernels.append(kernel_fn.kernel)	1✔
264
265	kernel_builder = KernelLaunchIRBuilder(lowerer.context, lowerer.builder)	1✔
266
267	ptr_to_queue_ref = kernel_builder.get_queue(exec_queue=kernel_fn.queue)	1✔
268
269	args = self._build_kernel_arglist(kernel_fn, lowerer, kernel_builder)	1✔
270	# Create a global range over which to submit the kernel based on the
271	# loop_ranges of the parfor
272	global_range = []	1✔
273
274	stop = _load_range(lowerer, 1)	1✔
275
276	global_range.append(stop)	1✔
277
278	local_range = []	1✔
279
280	kernel_ref_addr = kernel_fn.kernel.addressof_ref()	1✔
281	kernel_ref = lowerer.builder.inttoptr(	1✔
282	lowerer.context.get_constant(types.uintp, kernel_ref_addr),
283	cgutils.voidptr_t,
284	)
285	curr_queue_ref = lowerer.builder.load(ptr_to_queue_ref)	1✔
286
287	# Submit a synchronous kernel
288	kernel_builder.submit_sycl_kernel(	1✔
289	sycl_kernel_ref=kernel_ref,
290	sycl_queue_ref=curr_queue_ref,
291	total_kernel_args=args.num_flattened_args,
292	arg_list=args.arg_vals,
293	arg_ty_list=args.arg_types,
294	global_range=global_range,
295	local_range=local_range,
296	)
297
298	# At this point we can free the DPCTLSyclQueueRef (curr_queue)
299	kernel_builder.free_queue(ptr_to_sycl_queue_ref=ptr_to_queue_ref)	1✔
300
301	def _reduction_codegen(	1✔
302	self,
303	parfor,
304	typemap,
305	nredvars,
306	parfor_redvars,
307	parfor_reddict,
308	lowerer,
309	parfor_output_arrays,
310	loop_ranges,
311	flags,
312	alias_map,
313	):
314	"""
315	Reduction kernel generation and submission.
316	"""
317	inputArrayName = None	1✔
318	for para in parfor.params:	1!
319	if not isinstance(typemap[para], DpnpNdArray):	1!
UNCOV 320	continue	×
321	inputArrayName = para	1✔
322	break	1✔
323
324	if inputArrayName is None:	1!
325	raise AssertionError	×
326
327	inputArrayType = typemap[inputArrayName]	1✔
328
329	reductionHelperList = []	1✔
330	for i in range(nredvars):	1✔
331	reductionHelper = ReductionHelper()	1✔
332	reductionHelper._allocate_partial_reduction_arrays(	1✔
333	parfor,
334	lowerer,
335	parfor_redvars[i],
336	inputArrayType,
337	)
338	reductionHelperList.append(reductionHelper)	1✔
339
340	reductionKernelVar = ReductionKernelVariables(	1✔
341	lowerer=lowerer,
342	parfor_node=parfor,
343	typemap=typemap,
344	parfor_outputs=parfor_output_arrays,
345	reductionHelperList=reductionHelperList,
346	)
347
348	if len(parfor.loop_nests) > 1:	1!
349	raise UnsupportedParforError(	×
350	"Reduction with nested loop is not yet supported."
351	)
352
353	parfor_kernel = create_reduction_main_kernel_for_parfor(	1✔
354	loop_ranges,
355	parfor,
356	typemap,
357	flags,
358	bool(alias_map),
359	reductionKernelVar,
360	parfor_reddict,
361	)
362
363	self._submit_reduction_main_parfor_kernel(	1✔
364	lowerer,
365	parfor_kernel,
366	reductionHelperList[0],
367	)
368
369	parfor_kernel = create_reduction_remainder_kernel_for_parfor(	1✔
370	parfor,
371	typemap,
372	flags,
373	bool(alias_map),
374	reductionKernelVar,
375	parfor_reddict,
376	reductionHelperList,
377	)
378
379	self._submit_reduction_remainder_parfor_kernel(	1✔
380	lowerer,
381	parfor_kernel,
382	)
383
384	reductionKernelVar.copy_final_sum_to_host(parfor_kernel)	1✔
385
386	def _lower_parfor_as_kernel(self, lowerer, parfor):	1✔
387	"""Lowers a parfor node created by the dpjit compiler to a
388	``numba_dpex.kernel``.
389
390	The general approach is as follows:
391
392	- The code from the parfor's init block is lowered normally
393	in the context of the current function.
394	- The body of the parfor is transformed into a kernel function.
395	- Dpctl runtime calls to submit the kernel are added.
396
397	"""
398	# We copy the typemap here because for race condition variable we'll
399	# update their type to array so they can be updated by the kernel.
400	orig_typemap = lowerer.fndesc.typemap	1✔
401
402	# replace original typemap with copy and restore the original at the
403	# end.
404	lowerer.fndesc.typemap = copy.copy(orig_typemap)	1✔
405
406	if config.DEBUG_ARRAY_OPT:	1!
407	print("lowerer.fndesc", lowerer.fndesc, type(lowerer.fndesc))	×
408
409	typemap = lowerer.fndesc.typemap	1✔
410	varmap = lowerer.varmap	1✔
411
412	loc = parfor.init_block.loc	1✔
413	scope = parfor.init_block.scope	1✔
414
415	# Lower the init block of the parfor.
416	for instr in parfor.init_block.body:	1✔
417	lowerer.lower_inst(instr)	1✔
418
419	for racevar in parfor.races:	1!
420	if racevar not in varmap:	×
421	rvtyp = typemap[racevar]	×
422	rv = ir.Var(scope, racevar, loc)	×
423	lowerer._alloca_var(rv.name, rvtyp)	×
424
425	alias_map = {}	1✔
426	arg_aliases = {}	1✔
427
428	find_potential_aliases_parfor(	1✔
429	parfor,
430	parfor.params,
431	typemap,
432	lowerer.func_ir,
433	alias_map,
434	arg_aliases,
435	)
436
437	# run get_parfor_outputs() and get_parfor_reductions() before
438	# kernel creation since Jumps are modified so CFG of loop_body
439	# dict will become invalid
440	if parfor.params is None:	1!
441	raise AssertionError	×
442
443	parfor_output_arrays = get_parfor_outputs(parfor, parfor.params)	1✔
444
445	# compile parfor body as a separate dpex kernel function
446	flags = copy.copy(parfor.flags)	1✔
447	flags.error_model = "numpy"	1✔
448
449	# Can't get here unless
450	# flags.set('auto_parallel', ParallelOptions(True)) # noqa: E800
451	index_var_typ = typemap[parfor.loop_nests[0].index_variable.name]	1✔
452
453	# index variables should have the same type, check rest of indices
454	for loop_nest in parfor.loop_nests[1:]:	1✔
455	if typemap[loop_nest.index_variable.name] != index_var_typ:	1!
456	raise AssertionError	×
457
458	loop_ranges = [	1✔
459	(loop_nest.start, loop_nest.stop, loop_nest.step)
460	for loop_nest in parfor.loop_nests
461	]
462
463	parfor_redvars, parfor_reddict = parfor.redvars, parfor.reddict	1✔
464
465	nredvars = len(parfor_redvars)	1✔
466	if nredvars > 0:	1✔
467	self._reduction_codegen(	1✔
468	parfor,
469	typemap,
470	nredvars,
471	parfor_redvars,
472	parfor_reddict,
473	lowerer,
474	parfor_output_arrays,
475	loop_ranges,
476	flags,
477	alias_map,
478	)
479	else:
480	try:	1✔
481	parfor_kernel = create_kernel_for_parfor(	1✔
482	lowerer,
483	parfor,
484	typemap,
485	flags,
486	loop_ranges,
487	bool(alias_map),
488	parfor.races,
489	parfor_output_arrays,
490	)
491	except Exception:	×
492	# FIXME: Make the exception more informative
493	raise UnsupportedParforError	×
494
495	# Finally submit the kernel
496	self._submit_parfor_kernel(	1✔
497	lowerer,
498	parfor_kernel,
499	loop_ranges,
500	)
501
502	# TODO: free the kernel at this point
503
504	# Restore the original typemap of the function that was replaced
505	# temporarily at the beginning of this function.
506	lowerer.fndesc.typemap = orig_typemap	1✔
507
508
509	class ParforLowerFactory:	1✔
510	"""A pseudo-factory class that maps a device filter string to a lowering
511	function.
512
513	Each Parfor instruction can have an optional "lowerer" attribute. The
514	lowerer attribute determines how the parfor instruction should be lowered
515	to LLVM IR. In addition, the lower attribute decides which parfor
516	instructions can be fused together.
517
518	The factory class maintains a dictionary mapping every device
519	type (filter string) encountered so far to a lowerer function for that
520	device type. At this point numba-dpex does not generate device-specific code
521	and the lowerer used is same for all device types. However, as a different
522	ParforLowerImpl instance is returned for every parfor instruction that has
523	a distinct compute-follows-data inferred device it prevents illegal
524	parfor fusion.
525	"""
526
527	device_to_lowerer_map = {}	1✔
528
529	@classmethod	1✔
530	def get_lowerer(cls, device):	1✔
531	try:	1✔
532	lowerer = ParforLowerFactory.device_to_lowerer_map[device]	1✔
533	except KeyError:	1✔
534	lowerer = ParforLowerImpl()._lower_parfor_as_kernel	1✔
535	ParforLowerFactory.device_to_lowerer_map[device] = lowerer	1✔
536
537	return lowerer	1✔

IntelPython / numba-dpex / 7187743552

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