5604690751

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

import copy

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

from numba_dpex.core import config
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,
)

from numba_dpex.core.datamodel.models import dpex_data_model_manager as dpex_dmm

# 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 _get_exec_queue(self, kernel_fn, lowerer):
        """Creates a stack variable storing the sycl queue pointer used to
        launch the kernel function.
        """
        self.kernel_builder = KernelLaunchIRBuilder(lowerer, kernel_fn.kernel)

        # Create a local variable storing a pointer to a DPCTLSyclQueueRef
        # pointer.
        self.curr_queue = self.kernel_builder.get_queue(
            exec_queue=kernel_fn.queue
        )

    def _build_kernel_arglist(self, kernel_fn, lowerer):
        """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
        self.args_list = self.kernel_builder.allocate_kernel_arg_array(
            num_flattened_args
        )
        self.args_ty_list = self.kernel_builder.allocate_kernel_arg_ty_array(
            num_flattened_args
        )
        # Populate the args_list and the args_ty_list LLVM arrays
        self.kernel_arg_num = 0
        for arg_num, arg in enumerate(kernel_fn.kernel_args):
            argtype = kernel_fn.kernel_arg_types[arg_num]
            llvm_val = _getvar(lowerer, arg)
            if isinstance(argtype, DpnpNdArray):
                datamodel = dpex_dmm.lookup(arg_type)
                self.kernel_builder.build_array_arg(
                    array_val=llvm_val,
                    array_data_model=datamodel,
                    array_rank=argtype.ndim,
                    arg_list=self.args_list,
                    args_ty_list=self.args_ty_list,
                    arg_num=self.kernel_arg_num,
                )
                self.kernel_arg_num += datamodel.flattened_field_count
            else:
                if argtype == types.complex64:
                    self.kernel_builder.build_complex_arg(
                        llvm_val,
                        types.float32,
                        self.args_list,
                        self.args_ty_list,
                        self.kernel_arg_num,
                    )
                    self.kernel_arg_num += 2
                elif argtype == types.complex128:
                    self.kernel_builder.build_complex_arg(
                        llvm_val,
                        types.float64,
                        self.args_list,
                        self.args_ty_list,
                        self.kernel_arg_num,
                    )
                    self.kernel_arg_num += 2
                else:
                    self.kernel_builder.build_arg(
                        llvm_val,
                        argtype,
                        self.args_list,
                        self.args_ty_list,
                        self.kernel_arg_num,
                    )
                    self.kernel_arg_num += 1

    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)

        self._get_exec_queue(kernel_fn, lowerer)
        self._build_kernel_arglist(kernel_fn, lowerer)
        # 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 = []

        # Submit a synchronous kernel
        self.kernel_builder.submit_sync_kernel(
            self.curr_queue,
            self.kernel_arg_num,
            self.args_list,
            self.args_ty_list,
            global_range,
            local_range,
        )

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

    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)

        self._get_exec_queue(kernel_fn, lowerer)
        self._build_kernel_arglist(kernel_fn, lowerer)
        # 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)
        )

        # Submit a synchronous kernel
        self.kernel_builder.submit_sync_kernel(
            self.curr_queue,
            self.kernel_arg_num,
            self.args_list,
            self.args_ty_list,
            global_range,
            local_range,
        )

    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)

        self._get_exec_queue(kernel_fn, lowerer)
        self._build_kernel_arglist(kernel_fn, lowerer)
        # 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 = []

        # Submit a synchronous kernel
        self.kernel_builder.submit_sync_kernel(
            self.curr_queue,
            self.kernel_arg_num,
            self.args_list,
            self.args_ty_list,
            global_range,
            local_range,
        )

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

IntelPython / numba-dpex / 5604690751

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