8380418912

Committed 21 Mar 2024 07:05PM UTC coverage: 79.874%. Remained the same

Build # 8380418912

Build Type

Pull #1400

github

Committed by

web-flow

Commit Message

Merge 737eebf7e into 611a59711

Pull Request Pull Request #1400: Feature/move experimental decorators to stable

Run Details

1531 of 2238 branches covered (68.41%)

Branch coverage included in aggregate %.

4 of 4 new or added lines in 3 files covered. (100.0%)

36 existing lines in 4 files now uncovered.

6442 of 7744 relevant lines covered (83.19%)

0.83 hits per line

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

73.42

/numba_dpex/core/parfors/kernel_builder.py

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

import copy
import sys
import warnings

import dpctl
import dpctl.program as dpctl_prog
from numba.core import ir, types
from numba.core.errors import NumbaParallelSafetyWarning
from numba.core.ir_utils import (
    add_offset_to_labels,
    get_name_var_table,
    get_unused_var_name,
    legalize_names,
    mk_unique_var,
    remove_dead,
    remove_dels,
    rename_labels,
    replace_var_names,
)
from numba.core.typing import signature
from numba.parfors import parfor

import numba_dpex as dpex
from numba_dpex.core import config

from ..descriptor import dpex_kernel_target
from ..types import DpnpNdArray, USMNdArray
from ..utils.kernel_templates import RangeKernelTemplate


class ParforKernel:
    def __init__(
        self,
        name,
        kernel,
        signature,
        kernel_args,
        kernel_arg_types,
        queue: dpctl.SyclQueue,
    ):
        self.name = name
        self.kernel = kernel
        self.signature = signature
        self.kernel_args = kernel_args
        self.kernel_arg_types = kernel_arg_types
        self.queue = queue


def _print_block(block):
    for i, inst in enumerate(block.body):
        print("    ", i, inst)


def _print_body(body_dict):
    """Pretty-print a set of IR blocks."""
    for label, block in body_dict.items():
        print("label: ", label)
        _print_block(block)


def _compile_kernel_parfor(
    sycl_queue, kernel_name, func_ir, argtypes, debug=False
):
    # Create a SPIRVKernel object
    kernel = dpex.core.kernel_interface.spirv_kernel.SpirvKernel(
        func_ir, kernel_name
    )

    # compile the kernel
    kernel.compile(
        args=argtypes,
        typing_ctx=dpex_kernel_target.typing_context,
        target_ctx=dpex_kernel_target.target_context,
        debug=debug,
        compile_flags=None,
    )

    dpctl_create_program_from_spirv_flags = []
    if debug or config.DPEX_OPT == 0:
        # if debug is ON we need to pass additional flags to igc.
        dpctl_create_program_from_spirv_flags = ["-g", "-cl-opt-disable"]

    # create a program
    kernel_bundle = dpctl_prog.create_program_from_spirv(
        sycl_queue,
        kernel.device_driver_ir_module,
        " ".join(dpctl_create_program_from_spirv_flags),
    )
    #  create a kernel
    sycl_kernel = kernel_bundle.get_sycl_kernel(kernel.module_name)

    return sycl_kernel


def _legalize_names_with_typemap(names, typemap):
    """Replace illegal characters in Numba IR var names.

    We use ir_utils.legalize_names to replace internal IR variable names
    containing illegal characters (e.g. period) with a legal character
    (underscore) so as to create legal variable names. The original variable
    names are in the typemap so we also need to add the legalized name to the
    typemap as well.
    """
    outdict = legalize_names(names)
    # For each pair in the dict of legalized names...
    for x, y in outdict.items():
        # If the name had some legalization change to it...
        if x != y:
            # Set the type of the new name the same as the type of the old name.
            typemap[y] = typemap[x]
    return outdict


def _to_scalar_from_0d(x):
    if isinstance(x, types.ArrayCompatible) and x.ndim == 0:
        return x.dtype
    else:
        return x


def _replace_var_with_array_in_block(vars, block, typemap, calltypes):
    new_block = []
    for inst in block.body:
        if isinstance(inst, ir.Assign) and inst.target.name in vars:
            const_node = ir.Const(0, inst.loc)
            const_var = ir.Var(
                inst.target.scope, mk_unique_var("$const_ind_0"), inst.loc
            )
            typemap[const_var.name] = types.uintp
            const_assign = ir.Assign(const_node, const_var, inst.loc)
            new_block.append(const_assign)

            setitem_node = ir.SetItem(
                inst.target, const_var, inst.value, inst.loc
            )
            calltypes[setitem_node] = signature(
                types.none,
                types.npytypes.Array(typemap[inst.target.name], 1, "C"),
                types.intp,
                typemap[inst.target.name],
            )
            new_block.append(setitem_node)
            continue
        elif isinstance(inst, parfor.Parfor):
            _replace_var_with_array_internal(
                vars, {0: inst.init_block}, typemap, calltypes
            )
            _replace_var_with_array_internal(
                vars, inst.loop_body, typemap, calltypes
            )

        new_block.append(inst)
    return new_block


def _replace_var_with_array_internal(vars, loop_body, typemap, calltypes):
    for label, block in loop_body.items():
        block.body = _replace_var_with_array_in_block(
            vars, block, typemap, calltypes
        )


def _replace_var_with_array(vars, loop_body, typemap, calltypes):
    _replace_var_with_array_internal(vars, loop_body, typemap, calltypes)
    for v in vars:
        el_typ = typemap[v]
        typemap.pop(v, None)
        typemap[v] = types.npytypes.Array(el_typ, 1, "C")


def _find_setitems_block(setitems, block, typemap):
    for inst in block.body:
        if isinstance(inst, ir.StaticSetItem) or isinstance(inst, ir.SetItem):
            setitems.add(inst.target.name)
        elif isinstance(inst, parfor.Parfor):
            _find_setitems_block(setitems, inst.init_block, typemap)
            _find_setitems_body(setitems, inst.loop_body, typemap)


def _find_setitems_body(setitems, loop_body, typemap):
    """
    Find the arrays that are written into (goes into setitems)
    """
    for label, block in loop_body.items():
        _find_setitems_block(setitems, block, typemap)


def _replace_sentinel_with_parfor_body(kernel_ir, sentinel_name, loop_body):
    # new label for splitting sentinel block
    new_label = max(loop_body.keys()) + 1

    # Search all the block in the kernel function for the sentinel assignment.
    for label, block in kernel_ir.blocks.items():
        for i, inst in enumerate(block.body):
            if (
                isinstance(inst, ir.Assign)
                and inst.target.name == sentinel_name
            ):
                # We found the sentinel assignment.
                loc = inst.loc
                scope = block.scope
                # split block across __sentinel__
                # A new block is allocated for the statements prior to the
                # sentinel but the new block maintains the current block label.
                prev_block = ir.Block(scope, loc)
                prev_block.body = block.body[:i]

                # The current block is used for statements after the sentinel.
                block.body = block.body[i + 1 :]  # noqa: E203
                # But the current block gets a new label.
                body_first_label = min(loop_body.keys())

                # The previous block jumps to the minimum labelled block of the
                # parfor body.
                prev_block.append(ir.Jump(body_first_label, loc))
                # Add all the parfor loop body blocks to the kernel function's
                # IR.
                for loop, b in loop_body.items():
                    kernel_ir.blocks[loop] = b
                body_last_label = max(loop_body.keys())
                kernel_ir.blocks[new_label] = block
                kernel_ir.blocks[label] = prev_block
                # Add a jump from the last parfor body block to the block
                # containing statements after the sentinel.
                kernel_ir.blocks[body_last_label].append(
                    ir.Jump(new_label, loc)
                )
                break
        else:
            continue
        break


def create_kernel_for_parfor(
    lowerer,
    parfor_node,
    typemap,
    flags,
    loop_ranges,
    has_aliases,
    races,
    parfor_outputs,
) -> ParforKernel:
    """
    Creates a numba_dpex.kernel function for a parfor node.

    There are two parts to this function:

        1) Code to iterate across the iteration space as defined by
           the schedule.
        2) The parfor body that does the work for a single point in
           the iteration space.

    Part 1 is created as Python text for simplicity with a sentinel
    assignment to mark the point in the IR where the parfor body
    should be added. This Python text is 'exec'ed into existence and its
    IR retrieved with run_frontend. The IR is scanned for the sentinel
    assignment where that basic block is split and the IR for the parfor
    body inserted.
    """
    loc = parfor_node.init_block.loc

    # The parfor body and the main function body share ir.Var nodes.
    # We have to do some replacements of Var names in the parfor body
    # to make them legal parameter names. If we don't copy then the
    # Vars in the main function also would incorrectly change their name.
    loop_body = copy.copy(parfor_node.loop_body)
    remove_dels(loop_body)

    parfor_dim = len(parfor_node.loop_nests)
    loop_indices = [
        loop_nest.index_variable.name for loop_nest in parfor_node.loop_nests
    ]

    # Get all the parfor params.
    parfor_params = parfor_node.params
    typemap = lowerer.fndesc.typemap

    # Compute just the parfor inputs as a set difference.
    parfor_inputs = sorted(list(set(parfor_params) - set(parfor_outputs)))

    for race in races:
        msg = (
            "Variable %s used in parallel loop may be written "
            "to simultaneously by multiple workers and may result "
            "in non-deterministic or unintended results." % race
        )
        warnings.warn(NumbaParallelSafetyWarning(msg, loc))

    _replace_var_with_array(races, loop_body, typemap, lowerer.fndesc.calltypes)

    # Reorder all the params so that inputs go first then outputs.
    parfor_params = parfor_inputs + parfor_outputs

    # Some Var and loop_indices may not have legal parameter names so create a
    # dict of potentially illegal param name to guaranteed legal name.
    param_dict = _legalize_names_with_typemap(parfor_params, typemap)
    ind_dict = _legalize_names_with_typemap(loop_indices, typemap)

    # Compute a new list of legal loop index names.
    legal_loop_indices = [ind_dict[v] for v in loop_indices]
    # Get the types of each parameter.
    param_types = [_to_scalar_from_0d(typemap[v]) for v in parfor_params]
    # Calculate types of args passed to the kernel function.
    func_arg_types = [typemap[v] for v in (parfor_inputs + parfor_outputs)]

    # Replace illegal parameter names in the loop body with legal ones.
    replace_var_names(loop_body, param_dict)

    # remember the name before legalizing as the actual arguments
    parfor_args = parfor_params
    # Change parfor_params to be legal names.
    parfor_params = [param_dict[v] for v in parfor_params]
    parfor_params_orig = parfor_params

    parfor_params = []
    ascontig = False
    for pindex in range(len(parfor_params_orig)):
        if (
            ascontig
            and pindex < len(parfor_inputs)
            and isinstance(param_types[pindex], types.npytypes.Array)
        ):
            parfor_params.append(parfor_params_orig[pindex] + "param")
        else:
            parfor_params.append(parfor_params_orig[pindex])

    # Change parfor body to replace illegal loop index vars with legal ones.
    replace_var_names(loop_body, ind_dict)

    loop_body_var_table = get_name_var_table(loop_body)
    sentinel_name = get_unused_var_name("__sentinel__", loop_body_var_table)

    if config.DEBUG_ARRAY_OPT >= 1:
        print("legal parfor_params = ", parfor_params, type(parfor_params))

    # Determine the unique names of the kernel functions.
    kernel_name = "__dpex_parfor_kernel_%s" % (parfor_node.id)

    kernel_template = RangeKernelTemplate(
        kernel_name=kernel_name,
        kernel_params=parfor_params,
        kernel_rank=parfor_dim,
        ivar_names=legal_loop_indices,
        sentinel_name=sentinel_name,
        loop_ranges=loop_ranges,
        param_dict=param_dict,
    )
    kernel_ir = kernel_template.kernel_ir

    if config.DEBUG_ARRAY_OPT:
        print("kernel_ir dump ", type(kernel_ir))
        kernel_ir.dump()
        print("loop_body dump ", type(loop_body))
        _print_body(loop_body)

    # rename all variables in kernel_ir afresh
    var_table = get_name_var_table(kernel_ir.blocks)
    new_var_dict = {}
    reserved_names = (
        [sentinel_name] + list(param_dict.values()) + legal_loop_indices
    )
    for name, var in var_table.items():
        if not (name in reserved_names):
            new_var_dict[name] = mk_unique_var(name)
    replace_var_names(kernel_ir.blocks, new_var_dict)
    if config.DEBUG_ARRAY_OPT:
        print("kernel_ir dump after renaming ")
        kernel_ir.dump()

    kernel_param_types = param_types

    if config.DEBUG_ARRAY_OPT:
        print(
            "kernel_param_types = ",
            type(kernel_param_types),
            "\n",
            kernel_param_types,
        )

    kernel_stub_last_label = max(kernel_ir.blocks.keys()) + 1

    # Add kernel stub last label to each parfor.loop_body label to prevent
    # label conflicts.
    loop_body = add_offset_to_labels(loop_body, kernel_stub_last_label)

    _replace_sentinel_with_parfor_body(kernel_ir, sentinel_name, loop_body)

    if config.DEBUG_ARRAY_OPT:
        print("kernel_ir last dump before renaming")
        kernel_ir.dump()

    kernel_ir.blocks = rename_labels(kernel_ir.blocks)
    remove_dels(kernel_ir.blocks)

    old_alias = flags.noalias
    if not has_aliases:
        if config.DEBUG_ARRAY_OPT:
            print("No aliases found so adding noalias flag.")
        flags.noalias = True

    remove_dead(kernel_ir.blocks, kernel_ir.arg_names, kernel_ir, typemap)

    if config.DEBUG_ARRAY_OPT:
        print("kernel_ir after remove dead")
        kernel_ir.dump()

    kernel_sig = signature(types.none, *kernel_param_types)

    if config.DEBUG_ARRAY_OPT:
        sys.stdout.flush()

    if config.DEBUG_ARRAY_OPT:
        print("after DUFunc inline".center(80, "-"))
        kernel_ir.dump()

    # The ParforLegalizeCFD pass has already ensured that the LHS and RHS
    # arrays are on same device. We can take the queue from the first input
    # array and use that to compile the kernel.

    exec_queue: dpctl.SyclQueue = None

    for arg in parfor_args:
        obj = typemap[arg]
        if isinstance(obj, DpnpNdArray):
            filter_string = obj.queue.sycl_device
            # FIXME: A better design is required so that we do not have to
            # create a queue every time.
            exec_queue = dpctl.get_device_cached_queue(filter_string)

    if not exec_queue:
        raise AssertionError(
            "No execution found for parfor. No way to compile the kernel!"
        )

    sycl_kernel = _compile_kernel_parfor(
        exec_queue,
        kernel_name,
        kernel_ir,
        kernel_param_types,
        debug=flags.debuginfo,
    )

    flags.noalias = old_alias

    if config.DEBUG_ARRAY_OPT:
        print("kernel_sig = ", kernel_sig)

    return ParforKernel(
        name=kernel_name,
        kernel=sycl_kernel,
        signature=kernel_sig,
        kernel_args=parfor_args,
        kernel_arg_types=func_arg_types,
        queue=exec_queue,
    )


def update_sentinel(kernel_ir, sentinel_name, kernel_body, new_label):
    """Searched all the blocks in the IR generated from a kernel template and
    replaces the __sentinel__ instruction with the actual op for the parfor.

    Args:
        kernel_ir : Numba FunctionIR that was generated from a kernel template
        sentinel_name : The name of the sentinel instruction that is to be
        replaced.
        kernel_body : The function body of the kernel template generated
        numba_dpex.kernel function
        new_label: The new label to be used for the basic block created to store
        the instructions that replaced the sentinel
    """
    for label, block in kernel_ir.blocks.items():
        for i, inst in enumerate(block.body):
            if (
                isinstance(inst, ir.Assign)
                and inst.target.name == sentinel_name
            ):
                # We found the sentinel assignment.
                loc = inst.loc
                scope = block.scope
                # split block across __sentinel__
                # A new block is allocated for the statements prior to the
                # sentinel but the new block maintains the current block label.
                prev_block = ir.Block(scope, loc)
                prev_block.body = block.body[:i]

                # The current block is used for statements after the sentinel.
                block.body = block.body[i + 1 :]  # noqa: E203
                # But the current block gets a new label.
                body_first_label = min(kernel_body.keys())

                # The previous block jumps to the minimum labelled block of the
                # parfor body.
                prev_block.append(ir.Jump(body_first_label, loc))

                # Add all the parfor loop body blocks to the kernel IR
                for loop, b in kernel_body.items():
                    kernel_ir.blocks[loop] = copy.copy(b)
                    kernel_ir.blocks[loop].body = copy.copy(
                        kernel_ir.blocks[loop].body
                    )

                body_last_label = max(kernel_body.keys())
                kernel_ir.blocks[new_label] = block
                kernel_ir.blocks[label] = prev_block
                # Add a jump from the last parfor body block to the block
                # containing statements after the sentinel.
                kernel_ir.blocks[body_last_label].append(
                    ir.Jump(new_label, loc)
                )

                break
        else:
            continue
        break

1	# SPDX-FileCopyrightText: 2020 - 2024 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import copy	1✔
6	import sys	1✔
7	import warnings	1✔
8
9	import dpctl	1✔
10	import dpctl.program as dpctl_prog	1✔
11	from numba.core import ir, types	1✔
12	from numba.core.errors import NumbaParallelSafetyWarning	1✔
13	from numba.core.ir_utils import (	1✔
14	add_offset_to_labels,
15	get_name_var_table,
16	get_unused_var_name,
17	legalize_names,
18	mk_unique_var,
19	remove_dead,
20	remove_dels,
21	rename_labels,
22	replace_var_names,
23	)
24	from numba.core.typing import signature	1✔
25	from numba.parfors import parfor	1✔
26
27	import numba_dpex as dpex	1✔
28	from numba_dpex.core import config	1✔
29
30	from ..descriptor import dpex_kernel_target	1✔
31	from ..types import DpnpNdArray, USMNdArray	1✔
32	from ..utils.kernel_templates import RangeKernelTemplate	1✔
33
34
35	class ParforKernel:	1✔
36	def __init__(	1✔
37	self,
38	name,
39	kernel,
40	signature,
41	kernel_args,
42	kernel_arg_types,
43	queue: dpctl.SyclQueue,
44	):
45	self.name = name	1✔
46	self.kernel = kernel	1✔
47	self.signature = signature	1✔
48	self.kernel_args = kernel_args	1✔
49	self.kernel_arg_types = kernel_arg_types	1✔
50	self.queue = queue	1✔
51
52
53	def _print_block(block):	1✔
UNCOV 54	for i, inst in enumerate(block.body):	×
55	print(" ", i, inst)	×
56
57
58	def _print_body(body_dict):	1✔
59	"""Pretty-print a set of IR blocks."""
UNCOV 60	for label, block in body_dict.items():	×
61	print("label: ", label)	×
62	_print_block(block)	×
63
64
65	def _compile_kernel_parfor(	1✔
66	sycl_queue, kernel_name, func_ir, argtypes, debug=False
67	):
68	# Create a SPIRVKernel object
69	kernel = dpex.core.kernel_interface.spirv_kernel.SpirvKernel(	1✔
70	func_ir, kernel_name
71	)
72
73	# compile the kernel
74	kernel.compile(	1✔
75	args=argtypes,
76	typing_ctx=dpex_kernel_target.typing_context,
77	target_ctx=dpex_kernel_target.target_context,
78	debug=debug,
79	compile_flags=None,
80	)
81
82	dpctl_create_program_from_spirv_flags = []	1✔
83	if debug or config.DPEX_OPT == 0:	1!
84	# if debug is ON we need to pass additional flags to igc.
UNCOV 85	dpctl_create_program_from_spirv_flags = ["-g", "-cl-opt-disable"]	×
86
87	# create a program
88	kernel_bundle = dpctl_prog.create_program_from_spirv(	1✔
89	sycl_queue,
90	kernel.device_driver_ir_module,
91	" ".join(dpctl_create_program_from_spirv_flags),
92	)
93	# create a kernel
94	sycl_kernel = kernel_bundle.get_sycl_kernel(kernel.module_name)	1✔
95
96	return sycl_kernel	1✔
97
98
99	def _legalize_names_with_typemap(names, typemap):	1✔
100	"""Replace illegal characters in Numba IR var names.
101
102	We use ir_utils.legalize_names to replace internal IR variable names
103	containing illegal characters (e.g. period) with a legal character
104	(underscore) so as to create legal variable names. The original variable
105	names are in the typemap so we also need to add the legalized name to the
106	typemap as well.
107	"""
108	outdict = legalize_names(names)	1✔
109	# For each pair in the dict of legalized names...
110	for x, y in outdict.items():	1✔
111	# If the name had some legalization change to it...
112	if x != y:	1✔
113	# Set the type of the new name the same as the type of the old name.
114	typemap[y] = typemap[x]	1✔
115	return outdict	1✔
116
117
118	def _to_scalar_from_0d(x):	1✔
119	if isinstance(x, types.ArrayCompatible) and x.ndim == 0:	1!
UNCOV 120	return x.dtype	×
121	else:
122	return x	1✔
123
124
125	def _replace_var_with_array_in_block(vars, block, typemap, calltypes):	1✔
126	new_block = []	1✔
127	for inst in block.body:	1✔
128	if isinstance(inst, ir.Assign) and inst.target.name in vars:	1!
129	const_node = ir.Const(0, inst.loc)	×
UNCOV 130	const_var = ir.Var(	×
131	inst.target.scope, mk_unique_var("$const_ind_0"), inst.loc
132	)
133	typemap[const_var.name] = types.uintp	×
134	const_assign = ir.Assign(const_node, const_var, inst.loc)	×
UNCOV 135	new_block.append(const_assign)	×
136
UNCOV 137	setitem_node = ir.SetItem(	×
138	inst.target, const_var, inst.value, inst.loc
139	)
UNCOV 140	calltypes[setitem_node] = signature(	×
141	types.none,
142	types.npytypes.Array(typemap[inst.target.name], 1, "C"),
143	types.intp,
144	typemap[inst.target.name],
145	)
146	new_block.append(setitem_node)	×
UNCOV 147	continue	×
148	elif isinstance(inst, parfor.Parfor):	1!
UNCOV 149	_replace_var_with_array_internal(	×
150	vars, {0: inst.init_block}, typemap, calltypes
151	)
UNCOV 152	_replace_var_with_array_internal(	×
153	vars, inst.loop_body, typemap, calltypes
154	)
155
156	new_block.append(inst)	1✔
157	return new_block	1✔
158
159
160	def _replace_var_with_array_internal(vars, loop_body, typemap, calltypes):	1✔
161	for label, block in loop_body.items():	1✔
162	block.body = _replace_var_with_array_in_block(	1✔
163	vars, block, typemap, calltypes
164	)
165
166
167	def _replace_var_with_array(vars, loop_body, typemap, calltypes):	1✔
168	_replace_var_with_array_internal(vars, loop_body, typemap, calltypes)	1✔
169	for v in vars:	1!
170	el_typ = typemap[v]	×
171	typemap.pop(v, None)	×
UNCOV 172	typemap[v] = types.npytypes.Array(el_typ, 1, "C")	×
173
174
175	def _find_setitems_block(setitems, block, typemap):	1✔
176	for inst in block.body:	×
177	if isinstance(inst, ir.StaticSetItem) or isinstance(inst, ir.SetItem):	×
178	setitems.add(inst.target.name)	×
179	elif isinstance(inst, parfor.Parfor):	×
180	_find_setitems_block(setitems, inst.init_block, typemap)	×
UNCOV 181	_find_setitems_body(setitems, inst.loop_body, typemap)	×
182
183
184	def _find_setitems_body(setitems, loop_body, typemap):	1✔
185	"""
186	Find the arrays that are written into (goes into setitems)
187	"""
188	for label, block in loop_body.items():	×
UNCOV 189	_find_setitems_block(setitems, block, typemap)	×
190
191
192	def _replace_sentinel_with_parfor_body(kernel_ir, sentinel_name, loop_body):	1✔
193	# new label for splitting sentinel block
194	new_label = max(loop_body.keys()) + 1	1✔
195
196	# Search all the block in the kernel function for the sentinel assignment.
197	for label, block in kernel_ir.blocks.items():	1!
198	for i, inst in enumerate(block.body):	1!
199	if (	1✔
200	isinstance(inst, ir.Assign)
201	and inst.target.name == sentinel_name
202	):
203	# We found the sentinel assignment.
204	loc = inst.loc	1✔
205	scope = block.scope	1✔
206	# split block across __sentinel__
207	# A new block is allocated for the statements prior to the
208	# sentinel but the new block maintains the current block label.
209	prev_block = ir.Block(scope, loc)	1✔
210	prev_block.body = block.body[:i]	1✔
211
212	# The current block is used for statements after the sentinel.
213	block.body = block.body[i + 1 :] # noqa: E203	1✔
214	# But the current block gets a new label.
215	body_first_label = min(loop_body.keys())	1✔
216
217	# The previous block jumps to the minimum labelled block of the
218	# parfor body.
219	prev_block.append(ir.Jump(body_first_label, loc))	1✔
220	# Add all the parfor loop body blocks to the kernel function's
221	# IR.
222	for loop, b in loop_body.items():	1✔
223	kernel_ir.blocks[loop] = b	1✔
224	body_last_label = max(loop_body.keys())	1✔
225	kernel_ir.blocks[new_label] = block	1✔
226	kernel_ir.blocks[label] = prev_block	1✔
227	# Add a jump from the last parfor body block to the block
228	# containing statements after the sentinel.
229	kernel_ir.blocks[body_last_label].append(	1✔
230	ir.Jump(new_label, loc)
231	)
232	break	1✔
233	else:
UNCOV 234	continue	×
235	break	1✔
236
237
238	def create_kernel_for_parfor(	1✔
239	lowerer,
240	parfor_node,
241	typemap,
242	flags,
243	loop_ranges,
244	has_aliases,
245	races,
246	parfor_outputs,
247	) -> ParforKernel:
248	"""
249	Creates a numba_dpex.kernel function for a parfor node.
250
251	There are two parts to this function:
252
253	1) Code to iterate across the iteration space as defined by
254	the schedule.
255	2) The parfor body that does the work for a single point in
256	the iteration space.
257
258	Part 1 is created as Python text for simplicity with a sentinel
259	assignment to mark the point in the IR where the parfor body
260	should be added. This Python text is 'exec'ed into existence and its
261	IR retrieved with run_frontend. The IR is scanned for the sentinel
262	assignment where that basic block is split and the IR for the parfor
263	body inserted.
264	"""
265	loc = parfor_node.init_block.loc	1✔
266
267	# The parfor body and the main function body share ir.Var nodes.
268	# We have to do some replacements of Var names in the parfor body
269	# to make them legal parameter names. If we don't copy then the
270	# Vars in the main function also would incorrectly change their name.
271	loop_body = copy.copy(parfor_node.loop_body)	1✔
272	remove_dels(loop_body)	1✔
273
274	parfor_dim = len(parfor_node.loop_nests)	1✔
275	loop_indices = [	1✔
276	loop_nest.index_variable.name for loop_nest in parfor_node.loop_nests
277	]
278
279	# Get all the parfor params.
280	parfor_params = parfor_node.params	1✔
281	typemap = lowerer.fndesc.typemap	1✔
282
283	# Compute just the parfor inputs as a set difference.
284	parfor_inputs = sorted(list(set(parfor_params) - set(parfor_outputs)))	1✔
285
286	for race in races:	1!
UNCOV 287	msg = (	×
288	"Variable %s used in parallel loop may be written "
289	"to simultaneously by multiple workers and may result "
290	"in non-deterministic or unintended results." % race
291	)
UNCOV 292	warnings.warn(NumbaParallelSafetyWarning(msg, loc))	×
293
294	_replace_var_with_array(races, loop_body, typemap, lowerer.fndesc.calltypes)	1✔
295
296	# Reorder all the params so that inputs go first then outputs.
297	parfor_params = parfor_inputs + parfor_outputs	1✔
298
299	# Some Var and loop_indices may not have legal parameter names so create a
300	# dict of potentially illegal param name to guaranteed legal name.
301	param_dict = _legalize_names_with_typemap(parfor_params, typemap)	1✔
302	ind_dict = _legalize_names_with_typemap(loop_indices, typemap)	1✔
303
304	# Compute a new list of legal loop index names.
305	legal_loop_indices = [ind_dict[v] for v in loop_indices]	1✔
306	# Get the types of each parameter.
307	param_types = [_to_scalar_from_0d(typemap[v]) for v in parfor_params]	1✔
308	# Calculate types of args passed to the kernel function.
309	func_arg_types = [typemap[v] for v in (parfor_inputs + parfor_outputs)]	1✔
310
311	# Replace illegal parameter names in the loop body with legal ones.
312	replace_var_names(loop_body, param_dict)	1✔
313
314	# remember the name before legalizing as the actual arguments
315	parfor_args = parfor_params	1✔
316	# Change parfor_params to be legal names.
317	parfor_params = [param_dict[v] for v in parfor_params]	1✔
318	parfor_params_orig = parfor_params	1✔
319
320	parfor_params = []	1✔
321	ascontig = False	1✔
322	for pindex in range(len(parfor_params_orig)):	1✔
323	if (	1!
324	ascontig
325	and pindex < len(parfor_inputs)
326	and isinstance(param_types[pindex], types.npytypes.Array)
327	):
UNCOV 328	parfor_params.append(parfor_params_orig[pindex] + "param")	×
329	else:
330	parfor_params.append(parfor_params_orig[pindex])	1✔
331
332	# Change parfor body to replace illegal loop index vars with legal ones.
333	replace_var_names(loop_body, ind_dict)	1✔
334
335	loop_body_var_table = get_name_var_table(loop_body)	1✔
336	sentinel_name = get_unused_var_name("__sentinel__", loop_body_var_table)	1✔
337
338	if config.DEBUG_ARRAY_OPT >= 1:	1!
UNCOV 339	print("legal parfor_params = ", parfor_params, type(parfor_params))	×
340
341	# Determine the unique names of the kernel functions.
342	kernel_name = "__dpex_parfor_kernel_%s" % (parfor_node.id)	1✔
343
344	kernel_template = RangeKernelTemplate(	1✔
345	kernel_name=kernel_name,
346	kernel_params=parfor_params,
347	kernel_rank=parfor_dim,
348	ivar_names=legal_loop_indices,
349	sentinel_name=sentinel_name,
350	loop_ranges=loop_ranges,
351	param_dict=param_dict,
352	)
353	kernel_ir = kernel_template.kernel_ir	1✔
354
355	if config.DEBUG_ARRAY_OPT:	1!
356	print("kernel_ir dump ", type(kernel_ir))	×
357	kernel_ir.dump()	×
358	print("loop_body dump ", type(loop_body))	×
UNCOV 359	_print_body(loop_body)	×
360
361	# rename all variables in kernel_ir afresh
362	var_table = get_name_var_table(kernel_ir.blocks)	1✔
363	new_var_dict = {}	1✔
364	reserved_names = (	1✔
365	[sentinel_name] + list(param_dict.values()) + legal_loop_indices
366	)
367	for name, var in var_table.items():	1✔
368	if not (name in reserved_names):	1✔
369	new_var_dict[name] = mk_unique_var(name)	1✔
370	replace_var_names(kernel_ir.blocks, new_var_dict)	1✔
371	if config.DEBUG_ARRAY_OPT:	1!
372	print("kernel_ir dump after renaming ")	×
UNCOV 373	kernel_ir.dump()	×
374
375	kernel_param_types = param_types	1✔
376
377	if config.DEBUG_ARRAY_OPT:	1!
UNCOV 378	print(	×
379	"kernel_param_types = ",
380	type(kernel_param_types),
381	"\n",
382	kernel_param_types,
383	)
384
385	kernel_stub_last_label = max(kernel_ir.blocks.keys()) + 1	1✔
386
387	# Add kernel stub last label to each parfor.loop_body label to prevent
388	# label conflicts.
389	loop_body = add_offset_to_labels(loop_body, kernel_stub_last_label)	1✔
390
391	_replace_sentinel_with_parfor_body(kernel_ir, sentinel_name, loop_body)	1✔
392
393	if config.DEBUG_ARRAY_OPT:	1!
394	print("kernel_ir last dump before renaming")	×
UNCOV 395	kernel_ir.dump()	×
396
397	kernel_ir.blocks = rename_labels(kernel_ir.blocks)	1✔
398	remove_dels(kernel_ir.blocks)	1✔
399
400	old_alias = flags.noalias	1✔
401	if not has_aliases:	1✔
402	if config.DEBUG_ARRAY_OPT:	1!
UNCOV 403	print("No aliases found so adding noalias flag.")	×
404	flags.noalias = True	1✔
405
406	remove_dead(kernel_ir.blocks, kernel_ir.arg_names, kernel_ir, typemap)	1✔
407
408	if config.DEBUG_ARRAY_OPT:	1!
409	print("kernel_ir after remove dead")	×
UNCOV 410	kernel_ir.dump()	×
411
412	kernel_sig = signature(types.none, *kernel_param_types)	1✔
413
414	if config.DEBUG_ARRAY_OPT:	1!
UNCOV 415	sys.stdout.flush()	×
416
417	if config.DEBUG_ARRAY_OPT:	1!
418	print("after DUFunc inline".center(80, "-"))	×
UNCOV 419	kernel_ir.dump()	×
420
421	# The ParforLegalizeCFD pass has already ensured that the LHS and RHS
422	# arrays are on same device. We can take the queue from the first input
423	# array and use that to compile the kernel.
424
425	exec_queue: dpctl.SyclQueue = None	1✔
426
427	for arg in parfor_args:	1✔
428	obj = typemap[arg]	1✔
429	if isinstance(obj, DpnpNdArray):	1✔
430	filter_string = obj.queue.sycl_device	1✔
431	# FIXME: A better design is required so that we do not have to
432	# create a queue every time.
433	exec_queue = dpctl.get_device_cached_queue(filter_string)	1✔
434
435	if not exec_queue:	1!
UNCOV 436	raise AssertionError(	×
437	"No execution found for parfor. No way to compile the kernel!"
438	)
439
440	sycl_kernel = _compile_kernel_parfor(	1✔
441	exec_queue,
442	kernel_name,
443	kernel_ir,
444	kernel_param_types,
445	debug=flags.debuginfo,
446	)
447
448	flags.noalias = old_alias	1✔
449
450	if config.DEBUG_ARRAY_OPT:	1!
UNCOV 451	print("kernel_sig = ", kernel_sig)	×
452
453	return ParforKernel(	1✔
454	name=kernel_name,
455	kernel=sycl_kernel,
456	signature=kernel_sig,
457	kernel_args=parfor_args,
458	kernel_arg_types=func_arg_types,
459	queue=exec_queue,
460	)
461
462
463	def update_sentinel(kernel_ir, sentinel_name, kernel_body, new_label):	1✔
464	"""Searched all the blocks in the IR generated from a kernel template and
465	replaces the __sentinel__ instruction with the actual op for the parfor.
466
467	Args:
468	kernel_ir : Numba FunctionIR that was generated from a kernel template
469	sentinel_name : The name of the sentinel instruction that is to be
470	replaced.
471	kernel_body : The function body of the kernel template generated
472	numba_dpex.kernel function
473	new_label: The new label to be used for the basic block created to store
474	the instructions that replaced the sentinel
475	"""
476	for label, block in kernel_ir.blocks.items():	1!
477	for i, inst in enumerate(block.body):	1✔
478	if (	1✔
479	isinstance(inst, ir.Assign)
480	and inst.target.name == sentinel_name
481	):
482	# We found the sentinel assignment.
483	loc = inst.loc	1✔
484	scope = block.scope	1✔
485	# split block across __sentinel__
486	# A new block is allocated for the statements prior to the
487	# sentinel but the new block maintains the current block label.
488	prev_block = ir.Block(scope, loc)	1✔
489	prev_block.body = block.body[:i]	1✔
490
491	# The current block is used for statements after the sentinel.
492	block.body = block.body[i + 1 :] # noqa: E203	1✔
493	# But the current block gets a new label.
494	body_first_label = min(kernel_body.keys())	1✔
495
496	# The previous block jumps to the minimum labelled block of the
497	# parfor body.
498	prev_block.append(ir.Jump(body_first_label, loc))	1✔
499
500	# Add all the parfor loop body blocks to the kernel IR
501	for loop, b in kernel_body.items():	1✔
502	kernel_ir.blocks[loop] = copy.copy(b)	1✔
503	kernel_ir.blocks[loop].body = copy.copy(	1✔
504	kernel_ir.blocks[loop].body
505	)
506
507	body_last_label = max(kernel_body.keys())	1✔
508	kernel_ir.blocks[new_label] = block	1✔
509	kernel_ir.blocks[label] = prev_block	1✔
510	# Add a jump from the last parfor body block to the block
511	# containing statements after the sentinel.
512	kernel_ir.blocks[body_last_label].append(	1✔
513	ir.Jump(new_label, loc)
514	)
515
516	break	1✔
517	else:
518	continue	1✔
519	break	1✔

IntelPython / numba-dpex / 8380418912

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