5049504418

Build Type

push

github

Committed by GitHub

Commit Message

Merge pull request #1030 from IntelPython/feature/support_numba057

Run Details

1089 of 1674 branches covered (65.05%)

Branch coverage included in aggregate %.

153 of 214 new or added lines in 15 files covered. (71.5%)

8 existing lines in 2 files now uncovered.

4475 of 5643 relevant lines covered (79.3%)

0.79 hits per line

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

87.39

/numba_dpex/core/parfors/parfor_lowerer.py

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

import copy

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.utils.kernel_launcher import KernelLaunchIRBuilder
from numba_dpex.core.parfors.reduction_helper import (
    ReductionHelper,
    ReductionKernelVariables,
)

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,
)

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

IntelPython / numba-dpex / 5049504418

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