Add cuda support for points, line, and area glyphs

datashader/__init__.py

@@ -10,13 +10,8 @@
 from .glyphs import Point                                # noqa (API import)
 from .pipeline import Pipeline                           # noqa (API import)
 from . import transfer_functions as tf                   # noqa (API import)
+from . import datastructures                             # noqa (API import)
 
-from . import pandas                         # noqa (build backend dispatch)
-from . import xarray                         # noqa (build backend dispatch)
-try:
-    from . import dask                       # noqa (build backend dispatch)
-except ImportError:
-    pass
 
 # Make RaggedArray pandas extension array available for
 # pandas >= 0.24.0 is installed

datashader/compiler.py

@@ -15,7 +15,7 @@
 
 
 @memoize
-def compile_components(agg, schema, glyph):
+def compile_components(agg, schema, glyph, cuda=False):
     """Given a ``Aggregation`` object and a schema, return 5 sub-functions.
 
     Parameters
@@ -54,13 +54,13 @@ def compile_components(agg, schema, glyph):
     bases = list(unique(concat(r._bases for r in reds)))
     dshapes = [b.out_dshape(schema) for b in bases]
     # List of tuples of (append, base, input columns, temps)
-    calls = [_get_call_tuples(b, d) for (b, d) in zip(bases, dshapes)]
+    calls = [_get_call_tuples(b, d, schema, cuda) for (b, d) in zip(bases, dshapes)]
     # List of unique column names needed
     cols = list(unique(concat(pluck(2, calls))))
     # List of temps needed
     temps = list(pluck(3, calls))
 
-    create = make_create(bases, dshapes)
+    create = make_create(bases, dshapes, cuda)
     info = make_info(cols)
     append = make_append(bases, cols, calls, glyph)
     combine = make_combine(bases, dshapes, temps)
@@ -79,13 +79,19 @@ def traverse_aggregation(agg):
         yield agg
 
 
-def _get_call_tuples(base, dshape):
-    return base._build_append(dshape), (base,), base.inputs, base._temps
+def _get_call_tuples(base, dshape, schema, use_cuda):
+    return base._build_append(dshape, schema, use_cuda), (base,), base.inputs, base._temps
 
 
-def make_create(bases, dshapes):
+def make_create(bases, dshapes, cuda):
     creators = [b._build_create(d) for (b, d) in zip(bases, dshapes)]
-    return lambda shape: tuple(c(shape) for c in creators)
+    if cuda:
+        import cupy
+        array_module = cupy
+    else:
+        array_module = np
+    # array_module = np
+    return lambda shape: tuple(c(shape, array_module) for c in creators)
 
 
 def make_info(cols):

datashader/core.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 from numbers import Number
+from math import log10
 
 import numpy as np
 import pandas as pd
@@ -18,6 +19,11 @@
 from .resampling import resample_2d, resample_2d_distributed
 from . import reductions as rd
 
+try:
+    import cudf
+except ImportError:
+    cudf = None
+
 
 class Axis(object):
     """Interface for implementing axis transformations.
@@ -106,7 +112,7 @@ class LogAxis(Axis):
     @staticmethod
     @ngjit
     def mapper(val):
-        return np.log10(val)
+        return log10(val)
 
     @staticmethod
     @ngjit
@@ -285,6 +291,10 @@ def line(self, source, x, y, agg=None, axis=0):
         x, y = _broadcast_column_specifications(x, y)
 
         if axis == 0:
+            if cudf and isinstance(source, cudf.DataFrame):
+                raise ValueError("""\
+Canvas.line using a cudf GPU DataFrame is only supported with axis=1""")
+
             if (isinstance(x, (Number, string_types)) and
                     isinstance(y, (Number, string_types))):
                 glyph = LineAxis0(x, y)
@@ -467,6 +477,10 @@ def area(self, source, x, y, agg=None, axis=0, y_stack=None):
         x, y, y_stack = _broadcast_column_specifications(x, y, y_stack)
 
         if axis == 0:
+            if cudf and isinstance(source, cudf.DataFrame):
+                raise ValueError("""\
+Canvas.area using a cudf GPU DataFrame is only supported with axis=1""")
+
             if y_stack is None:
                 if (isinstance(x, (Number, string_types)) and
                         isinstance(y, (Number, string_types))):
@@ -992,7 +1006,8 @@ def bypixel(source, canvas, glyph, agg):
         source = source.drop([col for col in columns if col not in cols_to_keep])
         source = source.to_dask_dataframe()
 
-    if isinstance(source, pd.DataFrame):
+    if (isinstance(source, pd.DataFrame) or
+            (cudf and isinstance(source, cudf.DataFrame))):
         # Avoid datashape.Categorical instantiation bottleneck
         # by only retaining the necessary columns:
         # https://github.com/bokeh/datashader/issues/396

datashader/datastructures/__init__.py

@@ -0,0 +1,14 @@
+from . import pandas, xarray  # noqa (API import)
+
+try:
+    import dask as _dask  # noqa (Test dask installed)
+    from . import dask    # noqa (API import)
+except ImportError:
+    pass
+
+try:
+    import cudf as _cudf  # noqa (Test cudf installed)
+    import cupy as _cupy  # noqa (Test cupy installed)
+    from . import cudf    # noqa (API import)
+except ImportError:
+    pass

datashader/datastructures/cudf.py

@@ -0,0 +1,9 @@
+from __future__ import absolute_import
+from datashader.datastructures.pandas import default
+from datashader.core import bypixel
+import cudf
+
+
+@bypixel.pipeline.register(cudf.DataFrame)
+def cudf_pipeline(df, schema, canvas, glyph, summary):
+    return default(glyph, df, schema, canvas, summary, cuda=True)

datashader/dask.py → datashader/datastructures/dask.py

@@ -6,11 +6,11 @@
 from collections import OrderedDict
 from dask.base import tokenize, compute
 
-from .core import bypixel
-from .compatibility import apply
-from .compiler import compile_components
-from .glyphs import Glyph, LineAxis0
-from .utils import Dispatcher
+from datashader.core import bypixel
+from datashader.compatibility import apply
+from datashader.compiler import compile_components
+from datashader.glyphs import Glyph, LineAxis0
+from datashader.utils import Dispatcher
 
 __all__ = ()
 

datashader/pandas.py → datashader/datastructures/pandas.py

@@ -2,11 +2,11 @@
 
 import pandas as pd
 
-from .core import bypixel
-from .compiler import compile_components
-from .glyphs.points import _PointLike
-from .glyphs.area import _AreaToLineLike
-from .utils import Dispatcher
+from datashader.core import bypixel
+from datashader.compiler import compile_components
+from datashader.glyphs.points import _PointLike
+from datashader.glyphs.area import _AreaToLineLike
+from datashader.utils import Dispatcher
 from collections import OrderedDict
 
 __all__ = ()
@@ -22,14 +22,14 @@ def pandas_pipeline(df, schema, canvas, glyph, summary):
 
 @glyph_dispatch.register(_PointLike)
 @glyph_dispatch.register(_AreaToLineLike)
-def default(glyph, df, schema, canvas, summary):
-    create, info, append, _, finalize = compile_components(summary, schema, glyph)
+def default(glyph, source, schema, canvas, summary, cuda=False):
+    create, info, append, _, finalize = compile_components(summary, schema, glyph, cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
     extend = glyph._build_extend(x_mapper, y_mapper, info, append)
 
-    x_range = canvas.x_range or glyph.compute_x_bounds(df)
-    y_range = canvas.y_range or glyph.compute_y_bounds(df)
+    x_range = canvas.x_range or glyph.compute_x_bounds(source)
+    y_range = canvas.y_range or glyph.compute_y_bounds(source)
 
     width = canvas.plot_width
     height = canvas.plot_height
@@ -41,7 +41,7 @@ def default(glyph, df, schema, canvas, summary):
     y_axis = canvas.y_axis.compute_index(y_st, height)
 
     bases = create((height, width))
-    extend(bases, df, x_st + y_st, x_range + y_range)
+    extend(bases, source, x_st + y_st, x_range + y_range)
 
     return finalize(bases,
                     coords=OrderedDict([(glyph.x_label, x_axis),

datashader/xarray.py → datashader/datastructures/xarray.py

@@ -1,9 +1,9 @@
 from __future__ import absolute_import
 from datashader.glyphs.quadmesh import _QuadMeshLike
-from datashader.pandas import default
-from .core import bypixel
+from datashader.datastructures.pandas import default
+from datashader.core import bypixel
 import xarray as xr
-from .utils import Dispatcher
+from datashader.utils import Dispatcher
 
 
 glyph_dispatch = Dispatcher()