Use CUDA mutexes to support CUDA max_n and min_n without dask

ianthomas23 · Mar 24, 2023 · Mar 9, 2023 · Mar 22, 2023 · Mar 20, 2023 · b37c4fa23db989e4898148ca3b319441c536f76a
commit b37c4fa23db989e4898148ca3b319441c536f76a
diff --git a/datashader/compiler.py b/datashader/compiler.py
@@ -76,7 +76,7 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
         self_intersect = False
         antialias_stage_2 = False
 
-    # List of tuples of (append, base, input columns, temps, combine temps)
+    # List of tuples of (append, base, input columns, temps, combine temps, uses cuda mutex)
     calls = [_get_call_tuples(b, d, schema, cuda, antialias, self_intersect)
              for (b, d) in zip(bases, dshapes)]
 
@@ -87,8 +87,8 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
     combine_temps = list(pluck(4, calls))
 
     create = make_create(bases, dshapes, cuda)
-    info = make_info(cols)
-    append = make_append(bases, cols, calls, glyph, isinstance(agg, by), antialias)
+    append, uses_cuda_mutex = make_append(bases, cols, calls, glyph, isinstance(agg, by), antialias)
+    info = make_info(cols, uses_cuda_mutex)
     combine = make_combine(bases, dshapes, temps, combine_temps, antialias, cuda)
     finalize = make_finalize(bases, agg, schema, cuda)
 
@@ -113,6 +113,7 @@ def _get_call_tuples(base, dshape, schema, cuda, antialias, self_intersect):
         base.inputs,  # cols
         base._build_temps(cuda),  # temps
         base._build_combine_temps(cuda),  # combine temps
+        cuda and base.uses_cuda_mutex(),  # uses cuda mutex
     )
 
 
@@ -126,13 +127,26 @@ def make_create(bases, dshapes, cuda):
     return lambda shape: tuple(c(shape, array_module) for c in creators)
 
 
-def make_info(cols):
-    return lambda df: tuple(c.apply(df) for c in cols)
+def make_info(cols, uses_cuda_mutex):
+    def info(df):
+        ret = tuple(c.apply(df) for c in cols)
+        if uses_cuda_mutex:
+            import cupy  # Guaranteed to be available if uses_cuda_mutex is True
+            mutex_array = cupy.zeros((1,), dtype=np.uint32)
+            ret += (mutex_array,)
+        return ret
+
+    return info
 
 
 def make_append(bases, cols, calls, glyph, categorical, antialias):
     names = ('_{0}'.format(i) for i in count())
     inputs = list(bases) + list(cols)
+    any_uses_cuda_mutex = any(call[5] for call in calls)
+    if any_uses_cuda_mutex:
+        # This adds an argument to the append() function that is the cuda mutex
+        # generated in make_info.
+        inputs += [True]  # Doesn't matter what is in the list.
     signature = [next(names) for i in inputs]
     arg_lk = dict(zip(inputs, signature))
     local_lk = {}
@@ -144,7 +158,7 @@ def make_append(bases, cols, calls, glyph, categorical, antialias):
     else:
         subscript = None
 
-    for func, bases, cols, temps, _ in calls:
+    for func, bases, cols, temps, _, uses_cuda_mutex in calls:
         local_lk.update(zip(temps, (next(names) for i in temps)))
         func_name = next(names)
         namespace[func_name] = func
@@ -164,6 +178,9 @@ def make_append(bases, cols, calls, glyph, categorical, antialias):
         if antialias:
             args.append("aa_factor")
 
+        if uses_cuda_mutex:
+            args.append(signature[-1])
+
         where_reduction = len(bases) == 1 and isinstance(bases[0], where)
         if where_reduction:
             update_index_arg_name = next(names)
@@ -200,7 +217,7 @@ def make_append(bases, cols, calls, glyph, categorical, antialias):
                 '    {2}'
                 ).format(subscript, ', '.join(signature), '\n    '.join(body))
     exec(code, namespace)
-    return ngjit(namespace['append'])
+    return ngjit(namespace['append']), any_uses_cuda_mutex
 
 
 def make_combine(bases, dshapes, temps, combine_temps, antialias, cuda):

diff --git a/datashader/reductions.py b/datashader/reductions.py
@@ -13,7 +13,7 @@
 
 try:
     from datashader.transfer_functions._cuda_utils import (
-        cuda_atomic_nanmin, cuda_atomic_nanmax, cuda_args)
+        cuda_atomic_nanmin, cuda_atomic_nanmax, cuda_args, cuda_mutex_lock, cuda_mutex_unlock)
 except ImportError:
     cuda_atomic_nanmin, cuda_atomic_nanmmax, cuda_args = None, None, None
 
@@ -245,6 +245,9 @@ class Reduction(Expr):
     def __init__(self, column=None):
         self.column = column
 
+    def uses_cuda_mutex(self):
+        return False
+
     def uses_row_index(self):
         return False
 
@@ -1324,6 +1327,9 @@ def _combine(aggs):
 
 
 class max_n(FloatingNReduction):
+    def uses_cuda_mutex(self):
+        return True
+
     def _antialias_stage_2(self, self_intersect, array_module):
         return ((AntialiasCombination.MAX, array_module.nan),)
 
@@ -1344,20 +1350,43 @@ def _append(x, y, agg, field):
                     return i
         return -1
 
-    def _build_append(self, dshape, schema, cuda, antialias, self_intersect):
-        if cuda:
-            raise ValueError("'max_n' reduction is not supported on the GPU")
-        return super()._build_append(dshape, schema, cuda, antialias, self_intersect)
+    # GPU append functions
+    @staticmethod
+    @nb_cuda.jit(device=True)
+    def _append_cuda(x, y, agg, field, mutex):
+        if not isnull(field):
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            n = agg.shape[2]
+            index = (x, y)
+            cuda_mutex_lock(mutex, index)
+            for i in range(n):
+                if isnull(agg[y, x, i]) or field > agg[y, x, i]:
+                    # Bump previous values along to make room for new value.
+                    for j in range(n-1, i, -1):
+                        agg[y, x, j] = agg[y, x, j-1]
+                    agg[y, x, i] = field
+
+                    cuda_mutex_unlock(mutex, index)
+                    return i
+            cuda_mutex_unlock(mutex, index)
+        return -1
 
     @staticmethod
     def _combine(aggs):
+        if cp is not None and isinstance(aggs[0], cp.ndarray):
+            raise NotImplementedError("max_n not supported on GPU with dask")
+
         ret = aggs[0]
         for i in range(1, len(aggs)):
             nanmax_n_in_place(ret, aggs[i])
         return ret
 
 
 class min_n(FloatingNReduction):
+    def uses_cuda_mutex(self):
+        return True
+
     def _antialias_requires_2_stages(self):
         return True
 
@@ -1381,13 +1410,33 @@ def _append(x, y, agg, field):
                     return i
         return -1
 
-    def _build_append(self, dshape, schema, cuda, antialias, self_intersect):
-        if cuda:
-            raise ValueError("'min_n' reduction is not supported on the GPU")
-        return super()._build_append(dshape, schema, cuda, antialias, self_intersect)
+    # GPU append functions
+    @staticmethod
+    @nb_cuda.jit(device=True)
+    def _append_cuda(x, y, agg, field, mutex):
+        if not isnull(field):
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            n = agg.shape[2]
+            index = (x, y)
+            cuda_mutex_lock(mutex, index)
+            for i in range(n):
+                if isnull(agg[y, x, i]) or field < agg[y, x, i]:
+                    # Bump previous values along to make room for new value.
+                    for j in range(n-1, i, -1):
+                        agg[y, x, j] = agg[y, x, j-1]
+                    agg[y, x, i] = field
+
+                    cuda_mutex_unlock(mutex, index)
+                    return i
+            cuda_mutex_unlock(mutex, index)
+        return -1
 
     @staticmethod
     def _combine(aggs):
+        if cp is not None and isinstance(aggs[0], cp.ndarray):
+            raise NotImplementedError("min_n not supported on GPU with dask")
+
         ret = aggs[0]
         for i in range(1, len(aggs)):
             nanmin_n_in_place(ret, aggs[i])
@@ -1472,6 +1521,9 @@ def out_dshape(self, input_dshape, antialias):
         else:
             return dshape(ct.float64)
 
+    def uses_cuda_mutex(self):
+        return self.selector.uses_cuda_mutex()
+
     def uses_row_index(self):
         return self.column is None
 

diff --git a/datashader/tests/test_pandas.py b/datashader/tests/test_pandas.py
@@ -210,7 +210,7 @@ def test_max(df):
     assert_eq_xr(c.points(df, 'x', 'y', ds.max('f64')), out)
 
 
-@pytest.mark.parametrize('df', [df_pd])
+@pytest.mark.parametrize('df', dfs_pd)
 def test_min_n(df):
     solution = np.array([[[-3, -1, 0, 4, nan, nan], [-13, -11, 10, 12, 14, nan]],
                          [[-9, -7, -5, 6, 8, nan], [-19, -17, -15, 16, 18, nan]]])
@@ -222,7 +222,7 @@ def test_min_n(df):
             assert_eq_ndarray(agg[:, :, 0].data, c.points(df, 'x', 'y', ds.min('plusminus')).data)
 
 
-@pytest.mark.parametrize('df', [df_pd])
+@pytest.mark.parametrize('df', dfs_pd)
 def test_max_n(df):
     solution = np.array([[[4, 0, -1, -3, nan, nan], [14, 12, 10, -11, -13, nan]],
                          [[8, 6, -5, -7, -9, nan], [18, 16, -15, -17, -19, nan]]])
@@ -2601,14 +2601,10 @@ def test_canvas_size():
     ds.first_n('f64', n=3),
     ds.last('f64'),
     ds.last_n('f64', n=3),
-    ds.max_n('f64', n=3),
-    ds.min_n('f64', n=3),
     ds.where(ds.first('f64')),
     ds.where(ds.first_n('f64', n=3)),
     ds.where(ds.last('f64')),
     ds.where(ds.last_n('f64', n=3)),
-    ds.where(ds.max_n('f64', n=3)),
-    ds.where(ds.min_n('f64', n=3)),
 ])
 def test_reduction_on_cuda_raises_error(reduction):
     with pytest.raises(ValueError, match="not supported on the GPU"):

diff --git a/datashader/transfer_functions/_cuda_utils.py b/datashader/transfer_functions/_cuda_utils.py
@@ -171,3 +171,16 @@ def interp2d_kernel(x, xp, fp, left, right, output_y):
 
             # Update output
             output_y[i, j] = y_interp
+
+
+@cuda.jit(device=True)
+def cuda_mutex_lock(mutex, index):
+    while cuda.atomic.compare_and_swap(mutex, 0, 1) != 0:
+          pass
+    cuda.threadfence()
+
+
+@cuda.jit(device=True)
+def cuda_mutex_unlock(mutex, index):
+    cuda.threadfence()
+    cuda.atomic.exch(mutex, 0, 0)
diff --git a/doc/reduction.csv b/doc/reduction.csv
@@ -1,16 +1,16 @@
 ,                                        CPU, CPU + Dask, GPU, GPU + Dask, Antialiasing, Within :class:`~datashader.reductions.where`
-:class:`~datashader.reductions.any`,     yes, yes,        yes, yes,        yes,          
+:class:`~datashader.reductions.any`,     yes, yes,        yes, yes,        yes,
 :class:`~datashader.reductions.by`,      yes, yes,        yes, yes,        yes,
 :class:`~datashader.reductions.count`,   yes, yes,        yes, yes,        yes,
 :class:`~datashader.reductions.first`,   yes, ,           ,    ,           yes,          yes
 :class:`~datashader.reductions.first_n`, yes, ,           ,    ,           ,             yes
 :class:`~datashader.reductions.last`,    yes, ,           ,    ,           yes,          yes
 :class:`~datashader.reductions.last_n`,  yes, ,           ,    ,           ,             yes
 :class:`~datashader.reductions.max`,     yes, yes,        yes, yes,        yes,          yes
-:class:`~datashader.reductions.max_n`,   yes, yes,        ,    ,           ,             yes
+:class:`~datashader.reductions.max_n`,   yes, yes,        yes, ,           ,             yes
 :class:`~datashader.reductions.mean`,    yes, yes,        yes, yes,        yes,
 :class:`~datashader.reductions.min`,     yes, yes,        yes, yes,        yes,          yes
-:class:`~datashader.reductions.min_n`,   yes, yes,        ,    ,           ,             yes
+:class:`~datashader.reductions.min_n`,   yes, yes,        yes, ,           ,             yes
 :class:`~datashader.reductions.std`,     yes, yes,        ,    ,           ,
 :class:`~datashader.reductions.sum`,     yes, yes,        yes, yes,        yes,
 :class:`~datashader.reductions.var`,     yes, yes,        ,    ,           ,