|
2 | 2 | # |
3 | 3 | # SPDX-License-Identifier: Apache-2.0 |
4 | 4 |
|
5 | | -import numpy |
6 | 5 | from llvmlite import ir as llvmir |
7 | 6 | from llvmlite.ir import Constant |
8 | 7 | from numba.core import cgutils |
9 | 8 | from numba.core import config as numba_config |
10 | | -from numba.core import ir, types |
11 | | -from numba.core.ir_utils import ( |
12 | | - convert_size_to_var, |
13 | | - get_np_ufunc_typ, |
14 | | - mk_unique_var, |
15 | | -) |
| 9 | +from numba.core import types |
16 | 10 | from numba.core.typing import signature |
17 | 11 | from numba.extending import intrinsic, overload_classmethod |
18 | 12 | from numba.np.arrayobj import ( |
|
21 | 15 | make_array, |
22 | 16 | populate_array, |
23 | 17 | ) |
24 | | -from numba.np.ufunc.dufunc import DUFunc |
25 | 18 |
|
26 | 19 | from numba_dpex.core.runtime import context as dpexrt |
27 | 20 | from numba_dpex.core.types import DpnpNdArray |
28 | 21 |
|
29 | | -# Numpy array constructors |
30 | | - |
31 | | - |
32 | | -def _is_ufunc(func): |
33 | | - return isinstance(func, (numpy.ufunc, DUFunc)) or hasattr( |
34 | | - func, "is_dpnp_ufunc" |
35 | | - ) |
36 | | - |
37 | | - |
38 | | -def _mk_alloc( |
39 | | - typingctx, typemap, calltypes, lhs, size_var, dtype, scope, loc, lhs_typ |
40 | | -): |
41 | | - """generate an array allocation with np.empty() and return list of nodes. |
42 | | - size_var can be an int variable or tuple of int variables. |
43 | | - lhs_typ is the type of the array being allocated. |
44 | | - """ |
45 | | - out = [] |
46 | | - ndims = 1 |
47 | | - size_typ = types.intp |
48 | | - if isinstance(size_var, tuple): |
49 | | - if len(size_var) == 1: |
50 | | - size_var = size_var[0] |
51 | | - size_var = convert_size_to_var(size_var, typemap, scope, loc, out) |
52 | | - else: |
53 | | - # tuple_var = build_tuple([size_var...]) |
54 | | - ndims = len(size_var) |
55 | | - tuple_var = ir.Var(scope, mk_unique_var("$tuple_var"), loc) |
56 | | - if typemap: |
57 | | - typemap[tuple_var.name] = types.containers.UniTuple( |
58 | | - types.intp, ndims |
59 | | - ) |
60 | | - # constant sizes need to be assigned to vars |
61 | | - new_sizes = [ |
62 | | - convert_size_to_var(s, typemap, scope, loc, out) |
63 | | - for s in size_var |
64 | | - ] |
65 | | - tuple_call = ir.Expr.build_tuple(new_sizes, loc) |
66 | | - tuple_assign = ir.Assign(tuple_call, tuple_var, loc) |
67 | | - out.append(tuple_assign) |
68 | | - size_var = tuple_var |
69 | | - size_typ = types.containers.UniTuple(types.intp, ndims) |
70 | | - |
71 | | - if hasattr(lhs_typ, "__allocate__"): |
72 | | - return lhs_typ.__allocate__( |
73 | | - typingctx, |
74 | | - typemap, |
75 | | - calltypes, |
76 | | - lhs, |
77 | | - size_var, |
78 | | - dtype, |
79 | | - scope, |
80 | | - loc, |
81 | | - lhs_typ, |
82 | | - size_typ, |
83 | | - out, |
84 | | - ) |
85 | | - |
86 | | - # g_np_var = Global(numpy) |
87 | | - g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc) |
88 | | - if typemap: |
89 | | - typemap[g_np_var.name] = types.misc.Module(numpy) |
90 | | - g_np = ir.Global("np", numpy, loc) |
91 | | - g_np_assign = ir.Assign(g_np, g_np_var, loc) |
92 | | - # attr call: empty_attr = getattr(g_np_var, empty) |
93 | | - empty_attr_call = ir.Expr.getattr(g_np_var, "empty", loc) |
94 | | - attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc) |
95 | | - if typemap: |
96 | | - typemap[attr_var.name] = get_np_ufunc_typ(numpy.empty) |
97 | | - attr_assign = ir.Assign(empty_attr_call, attr_var, loc) |
98 | | - # Assume str(dtype) returns a valid type |
99 | | - dtype_str = str(dtype) |
100 | | - # alloc call: lhs = empty_attr(size_var, typ_var) |
101 | | - typ_var = ir.Var(scope, mk_unique_var("$np_typ_var"), loc) |
102 | | - if typemap: |
103 | | - typemap[typ_var.name] = types.functions.NumberClass(dtype) |
104 | | - # If dtype is a datetime/timedelta with a unit, |
105 | | - # then it won't return a valid type and instead can be created |
106 | | - # with a string. i.e. "datetime64[ns]") |
107 | | - if ( |
108 | | - isinstance(dtype, (types.NPDatetime, types.NPTimedelta)) |
109 | | - and dtype.unit != "" |
110 | | - ): |
111 | | - typename_const = ir.Const(dtype_str, loc) |
112 | | - typ_var_assign = ir.Assign(typename_const, typ_var, loc) |
113 | | - else: |
114 | | - if dtype_str == "bool": |
115 | | - # empty doesn't like 'bool' sometimes (e.g. kmeans example) |
116 | | - dtype_str = "bool_" |
117 | | - np_typ_getattr = ir.Expr.getattr(g_np_var, dtype_str, loc) |
118 | | - typ_var_assign = ir.Assign(np_typ_getattr, typ_var, loc) |
119 | | - alloc_call = ir.Expr.call(attr_var, [size_var, typ_var], (), loc) |
120 | | - |
121 | | - if calltypes: |
122 | | - cac = typemap[attr_var.name].get_call_type( |
123 | | - typingctx, [size_typ, types.functions.NumberClass(dtype)], {} |
124 | | - ) |
125 | | - # By default, all calls to "empty" are typed as returning a standard |
126 | | - # NumPy ndarray. If we are allocating a ndarray subclass here then |
127 | | - # just change the return type to be that of the subclass. |
128 | | - cac._return_type = ( |
129 | | - lhs_typ.copy(layout="C") if lhs_typ.layout == "F" else lhs_typ |
130 | | - ) |
131 | | - calltypes[alloc_call] = cac |
132 | | - if lhs_typ.layout == "F": |
133 | | - empty_c_typ = lhs_typ.copy(layout="C") |
134 | | - empty_c_var = ir.Var(scope, mk_unique_var("$empty_c_var"), loc) |
135 | | - if typemap: |
136 | | - typemap[empty_c_var.name] = lhs_typ.copy(layout="C") |
137 | | - empty_c_assign = ir.Assign(alloc_call, empty_c_var, loc) |
138 | | - |
139 | | - # attr call: asfortranarray = getattr(g_np_var, asfortranarray) |
140 | | - asfortranarray_attr_call = ir.Expr.getattr( |
141 | | - g_np_var, "asfortranarray", loc |
142 | | - ) |
143 | | - afa_attr_var = ir.Var( |
144 | | - scope, mk_unique_var("$asfortran_array_attr"), loc |
145 | | - ) |
146 | | - if typemap: |
147 | | - typemap[afa_attr_var.name] = get_np_ufunc_typ(numpy.asfortranarray) |
148 | | - afa_attr_assign = ir.Assign(asfortranarray_attr_call, afa_attr_var, loc) |
149 | | - # call asfortranarray |
150 | | - asfortranarray_call = ir.Expr.call(afa_attr_var, [empty_c_var], (), loc) |
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 | | - |
174 | 22 |
|
175 | 23 | def _empty_nd_impl(context, builder, arrtype, shapes): |
176 | 24 | """Utility function used for allocating a new array during LLVM code |
|
0 commit comments