6344850074

Committed 28 Sep 2023 10:24PM UTC coverage: 76.426% (-0.5%) from 76.941%

Build # 6344850074

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1129 from IntelPython/feature/switch_to_numba_0_58

Switch to numba 0.58

Run Details

1068 of 1622 branches covered (0.0%)

Branch coverage included in aggregate %.

5 of 7 new or added lines in 1 file covered. (71.43%)

27 existing lines in 1 file now uncovered.

4453 of 5602 relevant lines covered (79.49%)

0.79 hits per line

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

0.0

/numba_dpex/numba_patches/patch_mk_alloc.py

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


def patch():
    """
    Patches the numba.core.ir_utils.mk_alloc function to support non-NumPy array
    types.

    The patch extends the ir_utils.mk_alloc function that is used by
    numba.parfors.parfor to allocate an empty temp array inside a parfor
    section. Without the patch Numba always allocates the temp using
    numpy.empty, with the patch the function checks if the array type to be
    allocated has a __allocate__ method, and if so uses that method to allocate
    the temporary array.

    TODO: Already merged into Numba main and should be removed when numba-dpex
    is ported to Numba 0.58
    """
    import numpy
    from numba.core import ir, ir_utils, types
    from numba.core.ir_utils import (
        convert_size_to_var,
        get_np_ufunc_typ,
        mk_unique_var,
    )
    from numba.parfors import parfor

    def _mk_alloc(
        typingctx, typemap, calltypes, lhs, size_var, dtype, scope, loc, lhs_typ
    ):
        """generate an array allocation with np.empty() and return list of
        nodes. size_var can be an int variable or tuple of int variables.
        lhs_typ is the type of the array being allocated.
        """
        out = []
        ndims = 1
        size_typ = types.intp
        if isinstance(size_var, tuple):
            if len(size_var) == 1:
                size_var = size_var[0]
                size_var = convert_size_to_var(
                    size_var, typemap, scope, loc, out
                )
            else:
                # tuple_var = build_tuple([size_var...])
                ndims = len(size_var)
                tuple_var = ir.Var(scope, mk_unique_var("$tuple_var"), loc)
                if typemap:
                    typemap[tuple_var.name] = types.containers.UniTuple(
                        types.intp, ndims
                    )
                # constant sizes need to be assigned to vars
                new_sizes = [
                    convert_size_to_var(s, typemap, scope, loc, out)
                    for s in size_var
                ]
                tuple_call = ir.Expr.build_tuple(new_sizes, loc)
                tuple_assign = ir.Assign(tuple_call, tuple_var, loc)
                out.append(tuple_assign)
                size_var = tuple_var
                size_typ = types.containers.UniTuple(types.intp, ndims)

        if hasattr(lhs_typ, "__allocate__"):
            return lhs_typ.__allocate__(
                typingctx,
                typemap,
                calltypes,
                lhs,
                size_var,
                dtype,
                scope,
                loc,
                lhs_typ,
                size_typ,
                out,
            )

        # g_np_var = Global(numpy)
        g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
        if typemap:
            typemap[g_np_var.name] = types.misc.Module(numpy)
        g_np = ir.Global("np", numpy, loc)
        g_np_assign = ir.Assign(g_np, g_np_var, loc)
        # attr call: empty_attr = getattr(g_np_var, empty)
        empty_attr_call = ir.Expr.getattr(g_np_var, "empty", loc)
        attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)
        if typemap:
            typemap[attr_var.name] = get_np_ufunc_typ(numpy.empty)
        attr_assign = ir.Assign(empty_attr_call, attr_var, loc)
        # Assume str(dtype) returns a valid type
        dtype_str = str(dtype)
        # alloc call: lhs = empty_attr(size_var, typ_var)
        typ_var = ir.Var(scope, mk_unique_var("$np_typ_var"), loc)
        if typemap:
            typemap[typ_var.name] = types.functions.NumberClass(dtype)
        # If dtype is a datetime/timedelta with a unit,
        # then it won't return a valid type and instead can be created
        # with a string. i.e. "datetime64[ns]")
        if (
            isinstance(dtype, (types.NPDatetime, types.NPTimedelta))
            and dtype.unit != ""
        ):
            typename_const = ir.Const(dtype_str, loc)
            typ_var_assign = ir.Assign(typename_const, typ_var, loc)
        else:
            if dtype_str == "bool":
                # empty doesn't like 'bool' sometimes (e.g. kmeans example)
                dtype_str = "bool_"
            np_typ_getattr = ir.Expr.getattr(g_np_var, dtype_str, loc)
            typ_var_assign = ir.Assign(np_typ_getattr, typ_var, loc)
        alloc_call = ir.Expr.call(attr_var, [size_var, typ_var], (), loc)

        if calltypes:
            cac = typemap[attr_var.name].get_call_type(
                typingctx, [size_typ, types.functions.NumberClass(dtype)], {}
            )
            # By default, all calls to "empty" are typed as returning a standard
            # NumPy ndarray.  If we are allocating a ndarray subclass here then
            # just change the return type to be that of the subclass.
            cac._return_type = (
                lhs_typ.copy(layout="C") if lhs_typ.layout == "F" else lhs_typ
            )
            calltypes[alloc_call] = cac
        if lhs_typ.layout == "F":
            empty_c_typ = lhs_typ.copy(layout="C")
            empty_c_var = ir.Var(scope, mk_unique_var("$empty_c_var"), loc)
            if typemap:
                typemap[empty_c_var.name] = lhs_typ.copy(layout="C")
            empty_c_assign = ir.Assign(alloc_call, empty_c_var, loc)

            # attr call: asfortranarray = getattr(g_np_var, asfortranarray)
            asfortranarray_attr_call = ir.Expr.getattr(
                g_np_var, "asfortranarray", loc
            )
            afa_attr_var = ir.Var(
                scope, mk_unique_var("$asfortran_array_attr"), loc
            )
            if typemap:
                typemap[afa_attr_var.name] = get_np_ufunc_typ(
                    numpy.asfortranarray
                )
            afa_attr_assign = ir.Assign(
                asfortranarray_attr_call, afa_attr_var, loc
            )
            # call asfortranarray
            asfortranarray_call = ir.Expr.call(
                afa_attr_var, [empty_c_var], (), loc
            )
            if calltypes:
                calltypes[asfortranarray_call] = typemap[
                    afa_attr_var.name
                ].get_call_type(typingctx, [empty_c_typ], {})

            asfortranarray_assign = ir.Assign(asfortranarray_call, lhs, loc)

            out.extend(
                [
                    g_np_assign,
                    attr_assign,
                    typ_var_assign,
                    empty_c_assign,
                    afa_attr_assign,
                    asfortranarray_assign,
                ]
            )
        else:
            alloc_assign = ir.Assign(alloc_call, lhs, loc)
            out.extend([g_np_assign, attr_assign, typ_var_assign, alloc_assign])

        return out

    ir_utils.mk_alloc = _mk_alloc
    parfor.mk_alloc = _mk_alloc

1	# SPDX-FileCopyrightText: 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5
UNCOV 6	def patch():	×
7	"""
8	Patches the numba.core.ir_utils.mk_alloc function to support non-NumPy array
9	types.
10
11	The patch extends the ir_utils.mk_alloc function that is used by
12	numba.parfors.parfor to allocate an empty temp array inside a parfor
13	section. Without the patch Numba always allocates the temp using
14	numpy.empty, with the patch the function checks if the array type to be
15	allocated has a __allocate__ method, and if so uses that method to allocate
16	the temporary array.
17
18	TODO: Already merged into Numba main and should be removed when numba-dpex
19	is ported to Numba 0.58
20	"""
UNCOV 21	import numpy	×
UNCOV 22	from numba.core import ir, ir_utils, types	×
UNCOV 23	from numba.core.ir_utils import (	×
24	convert_size_to_var,
25	get_np_ufunc_typ,
26	mk_unique_var,
27	)
UNCOV 28	from numba.parfors import parfor	×
29
UNCOV 30	def _mk_alloc(	×
31	typingctx, typemap, calltypes, lhs, size_var, dtype, scope, loc, lhs_typ
32	):
33	"""generate an array allocation with np.empty() and return list of
34	nodes. size_var can be an int variable or tuple of int variables.
35	lhs_typ is the type of the array being allocated.
36	"""
UNCOV 37	out = []	×
UNCOV 38	ndims = 1	×
UNCOV 39	size_typ = types.intp	×
UNCOV 40	if isinstance(size_var, tuple):	×
UNCOV 41	if len(size_var) == 1:	×
UNCOV 42	size_var = size_var[0]	×
UNCOV 43	size_var = convert_size_to_var(	×
44	size_var, typemap, scope, loc, out
45	)
46	else:
47	# tuple_var = build_tuple([size_var...])
UNCOV 48	ndims = len(size_var)	×
UNCOV 49	tuple_var = ir.Var(scope, mk_unique_var("$tuple_var"), loc)	×
UNCOV 50	if typemap:	×
UNCOV 51	typemap[tuple_var.name] = types.containers.UniTuple(	×
52	types.intp, ndims
53	)
54	# constant sizes need to be assigned to vars
UNCOV 55	new_sizes = [	×
56	convert_size_to_var(s, typemap, scope, loc, out)
57	for s in size_var
58	]
UNCOV 59	tuple_call = ir.Expr.build_tuple(new_sizes, loc)	×
UNCOV 60	tuple_assign = ir.Assign(tuple_call, tuple_var, loc)	×
UNCOV 61	out.append(tuple_assign)	×
UNCOV 62	size_var = tuple_var	×
UNCOV 63	size_typ = types.containers.UniTuple(types.intp, ndims)	×
64
UNCOV 65	if hasattr(lhs_typ, "__allocate__"):	×
UNCOV 66	return lhs_typ.__allocate__(	×
67	typingctx,
68	typemap,
69	calltypes,
70	lhs,
71	size_var,
72	dtype,
73	scope,
74	loc,
75	lhs_typ,
76	size_typ,
77	out,
78	)
79
80	# g_np_var = Global(numpy)
81	g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)	×
82	if typemap:	×
83	typemap[g_np_var.name] = types.misc.Module(numpy)	×
84	g_np = ir.Global("np", numpy, loc)	×
85	g_np_assign = ir.Assign(g_np, g_np_var, loc)	×
86	# attr call: empty_attr = getattr(g_np_var, empty)
87	empty_attr_call = ir.Expr.getattr(g_np_var, "empty", loc)	×
88	attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)	×
89	if typemap:	×
90	typemap[attr_var.name] = get_np_ufunc_typ(numpy.empty)	×
91	attr_assign = ir.Assign(empty_attr_call, attr_var, loc)	×
92	# Assume str(dtype) returns a valid type
93	dtype_str = str(dtype)	×
94	# alloc call: lhs = empty_attr(size_var, typ_var)
95	typ_var = ir.Var(scope, mk_unique_var("$np_typ_var"), loc)	×
96	if typemap:	×
97	typemap[typ_var.name] = types.functions.NumberClass(dtype)	×
98	# If dtype is a datetime/timedelta with a unit,
99	# then it won't return a valid type and instead can be created
100	# with a string. i.e. "datetime64[ns]")
101	if (	×
102	isinstance(dtype, (types.NPDatetime, types.NPTimedelta))
103	and dtype.unit != ""
104	):
105	typename_const = ir.Const(dtype_str, loc)	×
106	typ_var_assign = ir.Assign(typename_const, typ_var, loc)	×
107	else:
108	if dtype_str == "bool":	×
109	# empty doesn't like 'bool' sometimes (e.g. kmeans example)
110	dtype_str = "bool_"	×
111	np_typ_getattr = ir.Expr.getattr(g_np_var, dtype_str, loc)	×
112	typ_var_assign = ir.Assign(np_typ_getattr, typ_var, loc)	×
113	alloc_call = ir.Expr.call(attr_var, [size_var, typ_var], (), loc)	×
114
115	if calltypes:	×
116	cac = typemap[attr_var.name].get_call_type(	×
117	typingctx, [size_typ, types.functions.NumberClass(dtype)], {}
118	)
119	# By default, all calls to "empty" are typed as returning a standard
120	# NumPy ndarray. If we are allocating a ndarray subclass here then
121	# just change the return type to be that of the subclass.
122	cac._return_type = (	×
123	lhs_typ.copy(layout="C") if lhs_typ.layout == "F" else lhs_typ
124	)
125	calltypes[alloc_call] = cac	×
126	if lhs_typ.layout == "F":	×
127	empty_c_typ = lhs_typ.copy(layout="C")	×
128	empty_c_var = ir.Var(scope, mk_unique_var("$empty_c_var"), loc)	×
129	if typemap:	×
130	typemap[empty_c_var.name] = lhs_typ.copy(layout="C")	×
131	empty_c_assign = ir.Assign(alloc_call, empty_c_var, loc)	×
132
133	# attr call: asfortranarray = getattr(g_np_var, asfortranarray)
134	asfortranarray_attr_call = ir.Expr.getattr(	×
135	g_np_var, "asfortranarray", loc
136	)
137	afa_attr_var = ir.Var(	×
138	scope, mk_unique_var("$asfortran_array_attr"), loc
139	)
140	if typemap:	×
141	typemap[afa_attr_var.name] = get_np_ufunc_typ(	×
142	numpy.asfortranarray
143	)
144	afa_attr_assign = ir.Assign(	×
145	asfortranarray_attr_call, afa_attr_var, loc
146	)
147	# call asfortranarray
148	asfortranarray_call = ir.Expr.call(	×
149	afa_attr_var, [empty_c_var], (), loc
150	)
151	if calltypes:	×
152	calltypes[asfortranarray_call] = typemap[	×
153	afa_attr_var.name
154	].get_call_type(typingctx, [empty_c_typ], {})
155
156	asfortranarray_assign = ir.Assign(asfortranarray_call, lhs, loc)	×
157
158	out.extend(	×
159	[
160	g_np_assign,
161	attr_assign,
162	typ_var_assign,
163	empty_c_assign,
164	afa_attr_assign,
165	asfortranarray_assign,
166	]
167	)
168	else:
169	alloc_assign = ir.Assign(alloc_call, lhs, loc)	×
170	out.extend([g_np_assign, attr_assign, typ_var_assign, alloc_assign])	×
171
172	return out	×
173
UNCOV 174	ir_utils.mk_alloc = _mk_alloc	×
UNCOV 175	parfor.mk_alloc = _mk_alloc	×

IntelPython / numba-dpex / 6344850074

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