1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360 | from os.path import dirname
from warnings import warn
import numpy as np
from ..utils import first_dict_item, updiv, get_cuda_srcfile
from ..misc.utils import ConvolutionInfos
from ..cuda.processing import CudaProcessing
class Convolution:
"""
A class for performing convolution on GPU with CUDA, but without using
textures (unlike for example in ``silx.opencl.convolution``)
"""
def __init__(self, shape, kernel, axes=None, mode=None, extra_options=None, cuda_options=None):
"""
Constructor of Cuda Convolution.
Parameters
-----------
shape: tuple
Shape of the array.
kernel: array-like
Convolution kernel (1D, 2D or 3D).
axes: tuple, optional
Axes along which the convolution is performed,
for batched convolutions.
mode: str, optional
Boundary handling mode. Available modes are:
- "reflect": cba|abcd|dcb
- "nearest": aaa|abcd|ddd
- "wrap": bcd|abcd|abc
- "constant": 000|abcd|000
Default is "reflect".
extra_options: dict, optional
Advanced options (dict). Current options are:
- "allocate_input_array": True
- "allocate_output_array": True
- "allocate_tmp_array": True
- "sourcemodule_kwargs": {}
- "batch_along_flat_dims": True
"""
self.cuda = CudaProcessing(**(cuda_options or {}))
self._configure_extra_options(extra_options)
self._determine_use_case(shape, kernel, axes)
self._allocate_memory(mode)
self._init_kernels()
def _configure_extra_options(self, extra_options):
self.extra_options = {
"allocate_input_array": True,
"allocate_output_array": True,
"allocate_tmp_array": True,
"sourcemodule_kwargs": {},
"batch_along_flat_dims": True,
}
extra_opts = extra_options or {}
self.extra_options.update(extra_opts)
self.sourcemodule_kwargs = self.extra_options["sourcemodule_kwargs"]
def _get_dimensions(self, shape, kernel):
self.shape = shape
self.data_ndim = self._check_dimensions(shape=shape, name="Data")
self.kernel_ndim = self._check_dimensions(arr=kernel, name="Kernel")
Nx = shape[-1]
if self.data_ndim >= 2:
Ny = shape[-2]
else:
Ny = 1
if self.data_ndim >= 3:
Nz = shape[-3]
else:
Nz = 1
self.Nx = np.int32(Nx)
self.Ny = np.int32(Ny)
self.Nz = np.int32(Nz)
def _determine_use_case(self, shape, kernel, axes):
"""
Determine the convolution use case from the input/kernel shape, and axes.
"""
self._get_dimensions(shape, kernel)
if self.kernel_ndim > self.data_ndim:
raise ValueError("Kernel dimensions cannot exceed data dimensions")
data_ndim = self.data_ndim
kernel_ndim = self.kernel_ndim
self.kernel = kernel.astype("f")
convol_infos = ConvolutionInfos()
k = (data_ndim, kernel_ndim)
if k not in convol_infos.use_cases:
raise ValueError(
"Cannot find a use case for data ndim = %d and kernel ndim = %d" % (data_ndim, kernel_ndim)
)
possible_use_cases = convol_infos.use_cases[k]
# If some dimensions are "flat", make a batched convolution along them
# Ex. data_dim = (1, Nx) -> batched 1D convolution
if self.extra_options["batch_along_flat_dims"] and (1 in self.shape):
axes = tuple([curr_dim for numels, curr_dim in zip(self.shape, range(len(self.shape))) if numels != 1])
#
self.use_case_name = None
for uc_name, uc_params in possible_use_cases.items():
if axes in convol_infos.allowed_axes[uc_name]:
self.use_case_name = uc_name
self.use_case_desc = uc_params["name"]
self.use_case_kernels = uc_params["kernels"].copy()
if self.use_case_name is None:
raise ValueError(
"Cannot find a use case for data ndim = %d, kernel ndim = %d and axes=%s"
% (data_ndim, kernel_ndim, str(axes))
)
# TODO implement this use case
if self.use_case_name == "batched_separable_2D_1D_3D":
raise NotImplementedError("The use case %s is not implemented" % self.use_case_name)
#
self.axes = axes
# Replace "axes=None" with an actual value (except for ND-ND)
allowed_axes = convol_infos.allowed_axes[self.use_case_name]
if len(allowed_axes) > 1:
# The default choice might impact perfs
self.axes = allowed_axes[0] or allowed_axes[1]
self.separable = self.use_case_name.startswith("separable")
self.batched = self.use_case_name.startswith("batched")
def _allocate_memory(self, mode):
self.mode = mode or "reflect"
# The current implementation does not support kernel size bigger than data size,
# except for mode="nearest"
for i, dim_size in enumerate(self.shape):
if min(self.kernel.shape) > dim_size and i in self.axes:
warn(
"Warning: kernel support is too large for data dimension %d (%d). Forcing convolution mode to 'nearest'"
% (i, dim_size),
Warning,
)
self.mode = "nearest"
#
option_array_names = {
"allocate_input_array": "data_in",
"allocate_output_array": "data_out",
"allocate_tmp_array": "data_tmp",
}
# Nonseparable transforms do not need tmp array
if not (self.separable):
self.extra_options["allocate_tmp_array"] = False
# Allocate arrays
for option_name, array_name in option_array_names.items():
if self.extra_options[option_name]:
value = self.cuda.allocate_array("value", self.shape, np.float32)
else:
value = None
setattr(self, array_name, value)
if isinstance(self.kernel, np.ndarray):
self.d_kernel = self.cuda.to_device("d_kernel", self.kernel)
else:
if not (isinstance(self.kernel, self.cuda.array_class)):
raise TypeError("kernel must be either numpy array or cupy array")
self.d_kernel = self.kernel
self._old_input_ref = None
self._old_output_ref = None
self._c_modes_mapping = {
"periodic": 2,
"wrap": 2,
"nearest": 1,
"replicate": 1,
"reflect": 0,
"constant": 3,
}
mp = self._c_modes_mapping
if self.mode.lower() not in mp:
raise ValueError(
"""
Mode %s is not available. Available modes are:
%s
"""
% (self.mode, str(mp.keys()))
)
if self.mode.lower() == "constant":
raise NotImplementedError("mode='constant' is not implemented yet")
self._c_conv_mode = mp[self.mode]
def _init_kernels(self):
if self.kernel_ndim > 1: # noqa: SIM102
if np.abs(np.diff(self.kernel.shape)).max() > 0:
raise NotImplementedError("Non-separable convolution with non-square kernels is not implemented yet")
# Compile source module
fname = get_cuda_srcfile("convolution.cu")
nabu_cuda_dir = dirname(fname)
compile_options = (f"-DUSED_CONV_MODE={self._c_conv_mode}", f"-I{nabu_cuda_dir}")
self.sourcemodule_kwargs["options"] = compile_options
# Blocks, grid
self._block_size = {1: (32, 1, 1), 2: (32, 32, 1), 3: (16, 8, 8)}[self.data_ndim] # TODO tune
self._n_blocks = tuple([int(updiv(a, b)) for a, b in zip(self.shape[::-1], self._block_size)])
# Prepare cuda kernel calls
self.cuda_kernels = {}
for axis, kern_name in enumerate(self.use_case_kernels):
self.cuda_kernels[axis] = self.cuda.kernel(kern_name, filename=fname, **self.sourcemodule_kwargs)
# Cuda kernel arguments
kernel_args = [
None, # input placeholder
None, # output placeholder
self.d_kernel,
np.int32(self.kernel.shape[0]),
self.Nx,
self.Ny,
self.Nz,
]
if self.kernel_ndim == 2:
kernel_args.insert(3, np.int32(self.kernel.shape[1]))
if self.kernel_ndim == 3:
kernel_args.insert(3, np.int32(self.kernel.shape[2]))
kernel_args.insert(4, np.int32(self.kernel.shape[1]))
self.kernel_args = tuple(kernel_args)
self.kernel_kwargs = {"grid": self._n_blocks, "block": self._block_size}
# If self.data_tmp is allocated, separable transforms can be performed
# by a series of batched transforms, without any copy, by swapping refs.
self.swap_pattern = None
if self.separable:
if self.data_tmp is not None:
self.swap_pattern = {
2: [("data_in", "data_tmp"), ("data_tmp", "data_out")],
3: [
("data_in", "data_out"),
("data_out", "data_tmp"),
("data_tmp", "data_out"),
],
}
else:
raise NotImplementedError("For now, data_tmp has to be allocated")
def _get_swapped_arrays(self, i):
"""
Get the input and output arrays to use when using a "swap pattern".
Swapping refs enables to avoid copies between temp. array and output.
For example, a separable 2D->1D convolution on 2D data reads:
data_tmp = convol(data_input, kernel, axis=1) # step i=0
data_out = convol(data_tmp, kernel, axis=0) # step i=1
:param i: current step number of the separable convolution
"""
n_batchs = len(self.axes)
in_ref, out_ref = self.swap_pattern[n_batchs][i]
d_in = getattr(self, in_ref)
d_out = getattr(self, out_ref)
return d_in, d_out
def _configure_kernel_args(self, cuda_kernel_args, input_ref, output_ref):
# TODO more elegant
if input_ref is not None or output_ref is not None:
cuda_kernel_args = list(cuda_kernel_args)
if input_ref is not None:
cuda_kernel_args[0] = input_ref
if output_ref is not None:
cuda_kernel_args[1] = output_ref
cuda_kernel_args = tuple(cuda_kernel_args)
return cuda_kernel_args
@staticmethod
def _check_dimensions(arr=None, shape=None, name="", dim_min=1, dim_max=3):
if shape is not None:
ndim = len(shape)
elif arr is not None:
ndim = arr.ndim
else:
raise ValueError("Please provide either arr= or shape=")
if ndim < dim_min or ndim > dim_max:
raise ValueError("%s dimensions should be between %d and %d" % (name, dim_min, dim_max))
return ndim
def _check_array(self, arr):
if not (isinstance(arr, self.cuda.array_class) or isinstance(arr, np.ndarray)): # noqa: SIM101
raise TypeError("Expected either cupy or numpy.ndarray")
if arr.dtype != np.float32:
raise TypeError("Data must be float32")
if arr.shape != self.shape:
raise ValueError("Expected data shape = %s" % str(self.shape))
def _set_arrays(self, array, output=None):
# Either copy H->D or update references.
if isinstance(array, np.ndarray):
self.data_in[:] = array[:]
else:
self._old_input_ref = self.data_in
self.data_in = array
data_in_ref = self.data_in
if output is not None: # noqa: SIM102
if not (isinstance(output, np.ndarray)):
self._old_output_ref = self.data_out
self.data_out = output
# Update Cuda kernel arguments with new array references
self.kernel_args = self._configure_kernel_args(self.kernel_args, data_in_ref, self.data_out)
def _separable_convolution(self):
assert len(self.axes) == len(self.use_case_kernels)
# Separable: one kernel call per data dimension
for i, axis in enumerate(self.axes):
in_ref, out_ref = self._get_swapped_arrays(i)
self._batched_convolution(axis, input_ref=in_ref, output_ref=out_ref)
def _batched_convolution(self, axis, input_ref=None, output_ref=None):
# Batched: one kernel call in total
cuda_kernel = self.cuda_kernels[axis]
cuda_kernel_args = self._configure_kernel_args(self.kernel_args, input_ref, output_ref)
ev = cuda_kernel(*cuda_kernel_args, **self.kernel_kwargs)
return ev
def _nd_convolution(self):
assert len(self.use_case_kernels) == 1
cuda_kernel = first_dict_item(self.use_case_kernels)
ev = cuda_kernel(*self.kernel_args, **self.kernel_kwargs)
return ev
def _recover_arrays_references(self):
if self._old_input_ref is not None:
self.data_in = self._old_input_ref
self._old_input_ref = None
if self._old_output_ref is not None:
self.data_out = self._old_output_ref
self._old_output_ref = None
self.kernel_args = self._configure_kernel_args(self.kernel_args, self.data_in, self.data_out)
def _get_output(self, output):
if output is None:
res = self.data_out.get()
else:
res = output
if isinstance(output, np.ndarray):
output[:] = self.data_out[:]
self._recover_arrays_references()
return res
def convolve(self, array, output=None):
"""
Convolve an array with the class kernel.
:param array: Input array. Can be numpy.ndarray or cupy array
:param output: Output array. Can be numpy.ndarray or cupy array
"""
self._check_array(array)
self._set_arrays(array, output=output)
if self.axes is not None:
if self.separable:
self._separable_convolution()
elif self.batched:
assert len(self.axes) == 1
self._batched_convolution(self.axes[0])
# else: ND-ND convol
else:
# ND-ND convol
self._nd_convolution()
res = self._get_output(output)
return res
__call__ = convolve
|