3
3
from numbers import Number
4
4
from types import FunctionType as ftype , BuiltinFunctionType as bftype
5
5
import sys
6
- #import importlib
7
- #import functools
8
6
import inspect
7
+ import dpctl
8
+ from dpctl .memory import MemoryUSMShared
9
9
10
10
debug = False
11
11
12
+
12
13
def dprint (* args ):
13
14
if debug :
14
15
print (* args )
15
16
sys .stdout .flush ()
16
17
17
- import dpctl
18
- from dpctl .memory import MemoryUSMShared
19
18
20
19
functions_list = [o [0 ] for o in getmembers (np ) if isfunction (o [1 ]) or isbuiltin (o [1 ])]
21
20
class_list = [o for o in getmembers (np ) if isclass (o [1 ])]
22
21
23
22
array_interface_property = "__array_interface__"
23
+
24
+
24
25
def has_array_interface (x ):
25
26
return hasattr (x , array_interface_property )
26
27
28
+
27
29
class ndarray (np .ndarray ):
28
30
"""
29
31
numpy.ndarray subclass whose underlying memory buffer is allocated
30
32
with a foreign allocator.
31
33
"""
32
- def __new__ (subtype , shape ,
33
- dtype = float , buffer = None , offset = 0 ,
34
- strides = None , order = None ):
34
+
35
+ def __new__ (
36
+ subtype , shape , dtype = float , buffer = None , offset = 0 , strides = None , order = None
37
+ ):
35
38
# Create a new array.
36
39
if buffer is None :
37
40
dprint ("dparray::ndarray __new__ buffer None" )
38
41
nelems = np .prod (shape )
39
42
dt = np .dtype (dtype )
40
43
isz = dt .itemsize
41
- nbytes = int (isz * max (1 , nelems ))
44
+ nbytes = int (isz * max (1 , nelems ))
42
45
buf = MemoryUSMShared (nbytes )
43
46
new_obj = np .ndarray .__new__ (
44
- subtype , shape , dtype = dt ,
45
- buffer = buf , offset = 0 ,
46
- strides = strides , order = order )
47
+ subtype ,
48
+ shape ,
49
+ dtype = dt ,
50
+ buffer = buf ,
51
+ offset = 0 ,
52
+ strides = strides ,
53
+ order = order ,
54
+ )
47
55
if hasattr (new_obj , array_interface_property ):
48
56
dprint ("buffer None new_obj already has sycl_usm" )
49
57
else :
@@ -55,9 +63,14 @@ def __new__(subtype, shape,
55
63
dprint ("dparray::ndarray __new__ buffer" , array_interface_property )
56
64
# also check for array interface
57
65
new_obj = np .ndarray .__new__ (
58
- subtype , shape , dtype = dtype ,
59
- buffer = buffer , offset = offset ,
60
- strides = strides , order = order )
66
+ subtype ,
67
+ shape ,
68
+ dtype = dtype ,
69
+ buffer = buffer ,
70
+ offset = offset ,
71
+ strides = strides ,
72
+ order = order ,
73
+ )
61
74
if hasattr (new_obj , array_interface_property ):
62
75
dprint ("buffer None new_obj already has sycl_usm" )
63
76
else :
@@ -68,17 +81,26 @@ def __new__(subtype, shape,
68
81
dprint ("dparray::ndarray __new__ buffer not None and not sycl_usm" )
69
82
nelems = np .prod (shape )
70
83
# must copy
71
- ar = np .ndarray (shape ,
72
- dtype = dtype , buffer = buffer ,
73
- offset = offset , strides = strides ,
74
- order = order )
84
+ ar = np .ndarray (
85
+ shape ,
86
+ dtype = dtype ,
87
+ buffer = buffer ,
88
+ offset = offset ,
89
+ strides = strides ,
90
+ order = order ,
91
+ )
75
92
nbytes = int (ar .nbytes )
76
93
buf = MemoryUSMShared (nbytes )
77
94
new_obj = np .ndarray .__new__ (
78
- subtype , shape , dtype = dtype ,
79
- buffer = buf , offset = 0 ,
80
- strides = strides , order = order )
81
- np .copyto (new_obj , ar , casting = 'no' )
95
+ subtype ,
96
+ shape ,
97
+ dtype = dtype ,
98
+ buffer = buf ,
99
+ offset = 0 ,
100
+ strides = strides ,
101
+ order = order ,
102
+ )
103
+ np .copyto (new_obj , ar , casting = "no" )
82
104
if hasattr (new_obj , array_interface_property ):
83
105
dprint ("buffer None new_obj already has sycl_usm" )
84
106
else :
@@ -89,7 +111,8 @@ def __new__(subtype, shape,
89
111
def __array_finalize__ (self , obj ):
90
112
dprint ("__array_finalize__:" , obj , hex (id (obj )), type (obj ))
91
113
# When called from the explicit constructor, obj is None
92
- if obj is None : return
114
+ if obj is None :
115
+ return
93
116
# When called in new-from-template, `obj` is another instance of our own
94
117
# subclass, that we might use to update the new `self` instance.
95
118
# However, when called from view casting, `obj` can be an instance of any
@@ -105,7 +128,11 @@ def __array_finalize__(self, obj):
105
128
dprint ("external_allocator:" , hex (ea ), type (ea ))
106
129
dprint ("data:" , hex (d ), type (d ))
107
130
dppl_rt_allocator = numba .dppl ._dppl_rt .get_external_allocator ()
108
- dprint ("dppl external_allocator:" , hex (dppl_rt_allocator ), type (dppl_rt_allocator ))
131
+ dprint (
132
+ "dppl external_allocator:" ,
133
+ hex (dppl_rt_allocator ),
134
+ type (dppl_rt_allocator ),
135
+ )
109
136
dprint (dir (mobj ))
110
137
if ea == dppl_rt_allocator :
111
138
return
@@ -118,24 +145,26 @@ def __array_finalize__(self, obj):
118
145
if hasattr (obj , array_interface_property ):
119
146
return
120
147
ob = ob .base
121
-
148
+
122
149
# Just raise an exception since __array_ufunc__ makes all reasonable cases not
123
150
# need the code below.
124
- raise ValueError ("Non-USM allocated ndarray can not viewed as a USM-allocated one without a copy" )
125
-
151
+ raise ValueError (
152
+ "Non-USM allocated ndarray can not viewed as a USM-allocated one without a copy"
153
+ )
154
+
126
155
# Tell Numba to not treat this type just like a NumPy ndarray but to propagate its type.
127
156
# This way it will use the custom dparray allocator.
128
157
__numba_no_subtype_ndarray__ = True
129
158
130
159
# Convert to a NumPy ndarray.
131
160
def as_ndarray (self ):
132
- return np .copy (self )
161
+ return np .copy (self )
133
162
134
163
def __array__ (self ):
135
164
return self
136
165
137
166
def __array_ufunc__ (self , ufunc , method , * inputs , ** kwargs ):
138
- if method == ' __call__' :
167
+ if method == " __call__" :
139
168
N = None
140
169
scalars = []
141
170
typing = []
@@ -162,41 +191,43 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
162
191
# USM memory. However, if kwarg has dparray-typed out then
163
192
# array_ufunc is called recursively so we cast out as regular
164
193
# NumPy ndarray (having a USM data pointer).
165
- if kwargs .get (' out' , None ) is None :
194
+ if kwargs .get (" out" , None ) is None :
166
195
# maybe copy?
167
196
# deal with multiple returned arrays, so kwargs['out'] can be tuple
168
197
res_type = np .result_type (* typing )
169
198
out = empty (inputs [0 ].shape , dtype = res_type )
170
199
out_as_np = np .ndarray (out .shape , out .dtype , out )
171
- kwargs [' out' ] = out_as_np
200
+ kwargs [" out" ] = out_as_np
172
201
else :
173
202
# If they manually gave dparray as out kwarg then we have to also
174
203
# cast as regular NumPy ndarray to avoid recursion.
175
- if isinstance (kwargs [' out' ], ndarray ):
176
- out = kwargs [' out' ]
177
- kwargs [' out' ] = np .ndarray (out .shape , out .dtype , out )
204
+ if isinstance (kwargs [" out" ], ndarray ):
205
+ out = kwargs [" out" ]
206
+ kwargs [" out" ] = np .ndarray (out .shape , out .dtype , out )
178
207
else :
179
- out = kwargs [' out' ]
208
+ out = kwargs [" out" ]
180
209
ret = ufunc (* scalars , ** kwargs )
181
210
return out
182
211
else :
183
212
return NotImplemented
184
213
214
+
185
215
def isdef (x ):
186
216
try :
187
217
eval (x )
188
218
return True
189
219
except NameError :
190
220
return False
191
221
222
+
192
223
for c in class_list :
193
224
cname = c [0 ]
194
225
if isdef (cname ):
195
226
continue
196
227
# For now we do the simple thing and copy the types from NumPy module into dparray module.
197
228
new_func = "%s = np.%s" % (cname , cname )
198
229
try :
199
- the_code = compile (new_func , ' __init__' , ' exec' )
230
+ the_code = compile (new_func , " __init__" , " exec" )
200
231
exec (the_code )
201
232
except :
202
233
print ("Failed to exec type propagation" , cname )
@@ -209,16 +240,18 @@ def isdef(x):
209
240
for fname in functions_list :
210
241
if isdef (fname ):
211
242
continue
212
- new_func = "def %s(*args, **kwargs):\n " % fname
243
+ new_func = "def %s(*args, **kwargs):\n " % fname
213
244
new_func += " ret = np.%s(*args, **kwargs)\n " % fname
214
245
new_func += " if type(ret) == np.ndarray:\n "
215
246
new_func += " ret = ndarray(ret.shape, ret.dtype, ret)\n "
216
247
new_func += " return ret\n "
217
- the_code = compile (new_func , ' __init__' , ' exec' )
248
+ the_code = compile (new_func , " __init__" , " exec" )
218
249
exec (the_code )
219
250
251
+
220
252
def from_ndarray (x ):
221
253
return copy (x )
222
254
255
+
223
256
def as_ndarray (x ):
224
- return np .copy (x )
257
+ return np .copy (x )
0 commit comments