Skip to content

Commit 3d2f7df

Browse files
ianthomas23ianthomas23
authored
Support max_n and min_n reductions on GPU (#1196)
* Use CUDA mutexes to support CUDA max_n and min_n without dask * Support CUDA with dask
1 parent ed0d58e commit 3d2f7df

6 files changed

Lines changed: 207 additions & 29 deletions

File tree

‎datashader/compiler.py‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
7676
self_intersect = False
7777
antialias_stage_2 = False
7878

79-
# List of tuples of (append, base, input columns, temps, combine temps)
79+
# List of tuples of (append, base, input columns, temps, combine temps, uses cuda mutex)
8080
calls = [_get_call_tuples(b, d, schema, cuda, antialias, self_intersect)
8181
for (b, d) in zip(bases, dshapes)]
8282

@@ -87,8 +87,8 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
8787
combine_temps = list(pluck(4, calls))
8888

8989
create = make_create(bases, dshapes, cuda)
90-
info = make_info(cols)
91-
append = make_append(bases, cols, calls, glyph, isinstance(agg, by), antialias)
90+
append, uses_cuda_mutex = make_append(bases, cols, calls, glyph, isinstance(agg, by), antialias)
91+
info = make_info(cols, uses_cuda_mutex)
9292
combine = make_combine(bases, dshapes, temps, combine_temps, antialias, cuda)
9393
finalize = make_finalize(bases, agg, schema, cuda)
9494

@@ -113,6 +113,7 @@ def _get_call_tuples(base, dshape, schema, cuda, antialias, self_intersect):
113113
base.inputs, # cols
114114
base._build_temps(cuda), # temps
115115
base._build_combine_temps(cuda), # combine temps
116+
cuda and base.uses_cuda_mutex(), # uses cuda mutex
116117
)
117118

118119

@@ -126,13 +127,26 @@ def make_create(bases, dshapes, cuda):
126127
return lambda shape: tuple(c(shape, array_module) for c in creators)
127128

128129

129-
def make_info(cols):
130-
return lambda df: tuple(c.apply(df) for c in cols)
130+
def make_info(cols, uses_cuda_mutex):
131+
def info(df):
132+
ret = tuple(c.apply(df) for c in cols)
133+
if uses_cuda_mutex:
134+
import cupy # Guaranteed to be available if uses_cuda_mutex is True
135+
mutex_array = cupy.zeros((1,), dtype=np.uint32)
136+
ret += (mutex_array,)
137+
return ret
138+
139+
return info
131140

132141

133142
def make_append(bases, cols, calls, glyph, categorical, antialias):
134143
names = ('_{0}'.format(i) for i in count())
135144
inputs = list(bases) + list(cols)
145+
any_uses_cuda_mutex = any(call[5] for call in calls)
146+
if any_uses_cuda_mutex:
147+
# This adds an argument to the append() function that is the cuda mutex
148+
# generated in make_info.
149+
inputs += ["_cuda_mutex"]
136150
signature = [next(names) for i in inputs]
137151
arg_lk = dict(zip(inputs, signature))
138152
local_lk = {}
@@ -144,7 +158,7 @@ def make_append(bases, cols, calls, glyph, categorical, antialias):
144158
else:
145159
subscript = None
146160

147-
for func, bases, cols, temps, _ in calls:
161+
for func, bases, cols, temps, _, uses_cuda_mutex in calls:
148162
local_lk.update(zip(temps, (next(names) for i in temps)))
149163
func_name = next(names)
150164
namespace[func_name] = func
@@ -164,6 +178,9 @@ def make_append(bases, cols, calls, glyph, categorical, antialias):
164178
if antialias:
165179
args.append("aa_factor")
166180

181+
if uses_cuda_mutex:
182+
args.append(arg_lk["_cuda_mutex"])
183+
167184
where_reduction = len(bases) == 1 and isinstance(bases[0], where)
168185
if where_reduction:
169186
update_index_arg_name = next(names)
@@ -200,7 +217,7 @@ def make_append(bases, cols, calls, glyph, categorical, antialias):
200217
' {2}'
201218
).format(subscript, ', '.join(signature), '\n '.join(body))
202219
exec(code, namespace)
203-
return ngjit(namespace['append'])
220+
return ngjit(namespace['append']), any_uses_cuda_mutex
204221

205222

206223
def make_combine(bases, dshapes, temps, combine_temps, antialias, cuda):

‎datashader/reductions.py‎

Lines changed: 86 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313

1414
try:
1515
from datashader.transfer_functions._cuda_utils import (
16-
cuda_atomic_nanmin, cuda_atomic_nanmax, cuda_args)
16+
cuda_atomic_nanmin, cuda_atomic_nanmax, cuda_args, cuda_mutex_lock, cuda_mutex_unlock,
17+
cuda_nanmax_n_in_place, cuda_nanmin_n_in_place)
1718
except ImportError:
1819
cuda_atomic_nanmin, cuda_atomic_nanmmax, cuda_args = None, None, None
1920

@@ -245,6 +246,9 @@ class Reduction(Expr):
245246
def __init__(self, column=None):
246247
self.column = column
247248

249+
def uses_cuda_mutex(self):
250+
return False
251+
248252
def uses_row_index(self):
249253
return False
250254

@@ -1324,6 +1328,9 @@ def _combine(aggs):
13241328

13251329

13261330
class max_n(FloatingNReduction):
1331+
def uses_cuda_mutex(self):
1332+
return True
1333+
13271334
def _antialias_stage_2(self, self_intersect, array_module):
13281335
return ((AntialiasCombination.MAX, array_module.nan),)
13291336

@@ -1344,10 +1351,33 @@ def _append(x, y, agg, field):
13441351
return i
13451352
return -1
13461353

1347-
def _build_append(self, dshape, schema, cuda, antialias, self_intersect):
1354+
# GPU append functions
1355+
@staticmethod
1356+
@nb_cuda.jit(device=True)
1357+
def _append_cuda(x, y, agg, field, mutex):
1358+
if not isnull(field):
1359+
# Linear walk along stored values.
1360+
# Could do binary search instead but not expecting n to be large.
1361+
n = agg.shape[2]
1362+
index = (x, y)
1363+
cuda_mutex_lock(mutex, index)
1364+
for i in range(n):
1365+
if isnull(agg[y, x, i]) or field > agg[y, x, i]:
1366+
# Bump previous values along to make room for new value.
1367+
for j in range(n-1, i, -1):
1368+
agg[y, x, j] = agg[y, x, j-1]
1369+
agg[y, x, i] = field
1370+
1371+
cuda_mutex_unlock(mutex, index)
1372+
return i
1373+
cuda_mutex_unlock(mutex, index)
1374+
return -1
1375+
1376+
def _build_combine(self, dshape, antialias, cuda):
13481377
if cuda:
1349-
raise ValueError("'max_n' reduction is not supported on the GPU")
1350-
return super()._build_append(dshape, schema, cuda, antialias, self_intersect)
1378+
return self._combine_cuda
1379+
else:
1380+
return self._combine
13511381

13521382
@staticmethod
13531383
def _combine(aggs):
@@ -1356,8 +1386,19 @@ def _combine(aggs):
13561386
nanmax_n_in_place(ret, aggs[i])
13571387
return ret
13581388

1389+
@staticmethod
1390+
def _combine_cuda(aggs):
1391+
ret = aggs[0]
1392+
kernel_args = cuda_args(ret.shape[:2])
1393+
for i in range(1, len(aggs)):
1394+
cuda_nanmax_n_in_place[kernel_args](ret, aggs[i])
1395+
return ret
1396+
13591397

13601398
class min_n(FloatingNReduction):
1399+
def uses_cuda_mutex(self):
1400+
return True
1401+
13611402
def _antialias_requires_2_stages(self):
13621403
return True
13631404

@@ -1381,10 +1422,33 @@ def _append(x, y, agg, field):
13811422
return i
13821423
return -1
13831424

1384-
def _build_append(self, dshape, schema, cuda, antialias, self_intersect):
1425+
# GPU append functions
1426+
@staticmethod
1427+
@nb_cuda.jit(device=True)
1428+
def _append_cuda(x, y, agg, field, mutex):
1429+
if not isnull(field):
1430+
# Linear walk along stored values.
1431+
# Could do binary search instead but not expecting n to be large.
1432+
n = agg.shape[2]
1433+
index = (x, y)
1434+
cuda_mutex_lock(mutex, index)
1435+
for i in range(n):
1436+
if isnull(agg[y, x, i]) or field < agg[y, x, i]:
1437+
# Bump previous values along to make room for new value.
1438+
for j in range(n-1, i, -1):
1439+
agg[y, x, j] = agg[y, x, j-1]
1440+
agg[y, x, i] = field
1441+
1442+
cuda_mutex_unlock(mutex, index)
1443+
return i
1444+
cuda_mutex_unlock(mutex, index)
1445+
return -1
1446+
1447+
def _build_combine(self, dshape, antialias, cuda):
13851448
if cuda:
1386-
raise ValueError("'min_n' reduction is not supported on the GPU")
1387-
return super()._build_append(dshape, schema, cuda, antialias, self_intersect)
1449+
return self._combine_cuda
1450+
else:
1451+
return self._combine
13881452

13891453
@staticmethod
13901454
def _combine(aggs):
@@ -1393,6 +1457,14 @@ def _combine(aggs):
13931457
nanmin_n_in_place(ret, aggs[i])
13941458
return ret
13951459

1460+
@staticmethod
1461+
def _combine_cuda(aggs):
1462+
ret = aggs[0]
1463+
kernel_args = cuda_args(ret.shape[:2])
1464+
for i in range(1, len(aggs)):
1465+
cuda_nanmin_n_in_place[kernel_args](ret, aggs[i])
1466+
return ret
1467+
13961468

13971469
class mode(Reduction):
13981470
"""Mode (most common value) of all the values encountered in ``column``.
@@ -1472,6 +1544,9 @@ def out_dshape(self, input_dshape, antialias):
14721544
else:
14731545
return dshape(ct.float64)
14741546

1547+
def uses_cuda_mutex(self):
1548+
return self.selector.uses_cuda_mutex()
1549+
14751550
def uses_row_index(self):
14761551
return self.column is None
14771552

@@ -1543,6 +1618,10 @@ def _build_bases(self, cuda=False):
15431618
return self.selector._build_bases(cuda=cuda) + super()._build_bases(cuda=cuda)
15441619

15451620
def _build_combine(self, dshape, antialias, cuda):
1621+
if cuda and self.uses_cuda_mutex():
1622+
raise NotImplementedError(
1623+
"'where' reduction does not support a selector that uses a CUDA mutex such as 'max_n'")
1624+
15461625
# Does not support categorical reductions.
15471626
selector = self.selector
15481627
append = selector._append

‎datashader/tests/test_dask.py‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ def test_max(ddf):
180180
assert_eq_xr(c.points(ddf, 'x', 'y', ds.max('f64')), out)
181181

182182

183-
@pytest.mark.parametrize('ddf', [_ddf])
183+
@pytest.mark.parametrize('ddf', ddfs)
184184
def test_min_n(ddf):
185185
solution = np.array([[[-3, -1, 0, 4, nan, nan], [-13, -11, 10, 12, 14, nan]],
186186
[[-9, -7, -5, 6, 8, nan], [-19, -17, -15, 16, 18, nan]]])
@@ -192,7 +192,7 @@ def test_min_n(ddf):
192192
assert_eq_ndarray(agg[:, :, 0].data, c.points(ddf, 'x', 'y', ds.min('plusminus')).data)
193193

194194

195-
@pytest.mark.parametrize('ddf', [_ddf])
195+
@pytest.mark.parametrize('ddf', ddfs)
196196
def test_max_n(ddf):
197197
solution = np.array([[[4, 0, -1, -3, nan, nan], [14, 12, 10, -11, -13, nan]],
198198
[[8, 6, -5, -7, -9, nan], [18, 16, -15, -17, -19, nan]]])
@@ -1694,6 +1694,19 @@ def test_canvas_size():
16941694
ddf = dd.from_pandas(df, 1)
16951695

16961696
for cvs in cvs_list:
1697-
with pytest.raises(ValueError, match=msg):
1697+
with pytest.raises(ValueError, match=msg):
16981698
cvs.points(ddf, "x", "y", ds.mean("z"))
1699-
1699+
1700+
1701+
@pytest.mark.skipif(not test_gpu, reason="DATASHADER_TEST_GPU not set")
1702+
@pytest.mark.parametrize('reduction', [
1703+
ds.where(ds.first('f64')),
1704+
ds.where(ds.first_n('f64')),
1705+
ds.where(ds.last('f64')),
1706+
ds.where(ds.last_n('f64')),
1707+
ds.where(ds.max_n('f64', n=3)),
1708+
ds.where(ds.min_n('f64', n=3)),
1709+
])
1710+
def test_reduction_on_cuda_dask_raises_error(reduction):
1711+
with pytest.raises((NotImplementedError, ValueError)):
1712+
c.points(cudf_ddf, 'x', 'y', reduction)

‎datashader/tests/test_pandas.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -106,8 +106,12 @@ def assert_eq_xr(agg, b, close=False):
106106
def assert_eq_ndarray(data, b, close=False):
107107
"""Assert that two ndarrays are equal, handling the possibility that the
108108
ndarrays are of different types"""
109-
if cupy and isinstance(data, cupy.ndarray):
110-
data = cupy.asnumpy(data)
109+
if cupy:
110+
if isinstance(data, cupy.ndarray):
111+
data = cupy.asnumpy(data)
112+
if isinstance(b, cupy.ndarray):
113+
b = cupy.asnumpy(b)
114+
111115
if close:
112116
np.testing.assert_array_almost_equal(data, b, decimal=5)
113117
else:
@@ -210,7 +214,7 @@ def test_max(df):
210214
assert_eq_xr(c.points(df, 'x', 'y', ds.max('f64')), out)
211215

212216

213-
@pytest.mark.parametrize('df', [df_pd])
217+
@pytest.mark.parametrize('df', dfs_pd)
214218
def test_min_n(df):
215219
solution = np.array([[[-3, -1, 0, 4, nan, nan], [-13, -11, 10, 12, 14, nan]],
216220
[[-9, -7, -5, 6, 8, nan], [-19, -17, -15, 16, 18, nan]]])
@@ -222,7 +226,7 @@ def test_min_n(df):
222226
assert_eq_ndarray(agg[:, :, 0].data, c.points(df, 'x', 'y', ds.min('plusminus')).data)
223227

224228

225-
@pytest.mark.parametrize('df', [df_pd])
229+
@pytest.mark.parametrize('df', dfs_pd)
226230
def test_max_n(df):
227231
solution = np.array([[[4, 0, -1, -3, nan, nan], [14, 12, 10, -11, -13, nan]],
228232
[[8, 6, -5, -7, -9, nan], [18, 16, -15, -17, -19, nan]]])
@@ -2601,14 +2605,10 @@ def test_canvas_size():
26012605
ds.first_n('f64', n=3),
26022606
ds.last('f64'),
26032607
ds.last_n('f64', n=3),
2604-
ds.max_n('f64', n=3),
2605-
ds.min_n('f64', n=3),
26062608
ds.where(ds.first('f64')),
26072609
ds.where(ds.first_n('f64', n=3)),
26082610
ds.where(ds.last('f64')),
26092611
ds.where(ds.last_n('f64', n=3)),
2610-
ds.where(ds.max_n('f64', n=3)),
2611-
ds.where(ds.min_n('f64', n=3)),
26122612
])
26132613
def test_reduction_on_cuda_raises_error(reduction):
26142614
with pytest.raises(ValueError, match="not supported on the GPU"):

0 commit comments

Comments
 (0)