Pre-compile antialias stage 2 combination (#1258)

* Move second-stage antialiasing to compiler.py

* Remove unnecessary use_2_stage_agg argument to line._build_extend functions

* Better function naming

datashader/antialias.py

@@ -2,11 +2,6 @@
 from enum import Enum
 from typing import NamedTuple, TYPE_CHECKING
 
-from datashader.utils import (
-    nanfirst_in_place, nanlast_in_place, nanmax_in_place,
-    nanmin_in_place, nansum_in_place, ngjit, parallel_fill)
-from numba import literal_unroll
-
 
 # Enum used to specify how the second stage aggregation is performed
 # for 2-stage antialiased lines.
@@ -23,10 +18,12 @@ class AntialiasStage2(NamedTuple):
     """Configuration for second-stage combination of a single antialiased reduction."""
     combination: AntialiasCombination
     zero: float
+    n_reduction: bool = False
+    categorical: bool = False
 
 
 if TYPE_CHECKING:
-    UnzippedAntialiasStage2 = tuple[tuple[AntialiasCombination], tuple[float]]
+    UnzippedAntialiasStage2 = tuple[tuple[AntialiasCombination], tuple[float], tuple[bool], tuple[bool]]
 
 
 def two_stage_agg(antialias_stage_2: UnzippedAntialiasStage2 | None):
@@ -56,42 +53,3 @@ def two_stage_agg(antialias_stage_2: UnzippedAntialiasStage2 | None):
             break
 
     return overwrite, use_2_stage_agg
-
-
-@ngjit
-def _combine_in_place(accum_agg, other_agg, antialias_combination):
-    if antialias_combination == AntialiasCombination.MAX:
-        nanmax_in_place(accum_agg, other_agg)
-    elif antialias_combination == AntialiasCombination.MIN:
-        nanmin_in_place(accum_agg, other_agg)
-    elif antialias_combination == AntialiasCombination.FIRST:
-        nanfirst_in_place(accum_agg, other_agg)
-    elif antialias_combination == AntialiasCombination.LAST:
-        nanlast_in_place(accum_agg, other_agg)
-    else:
-        nansum_in_place(accum_agg, other_agg)
-
-
-@ngjit
-def aa_stage_2_accumulate(aggs_and_copies, antialias_combinations):
-    k = 0
-    # Numba access to heterogeneous tuples is only permitted using literal_unroll.
-    for agg_and_copy in literal_unroll(aggs_and_copies):
-        _combine_in_place(agg_and_copy[1], agg_and_copy[0], antialias_combinations[k])
-        k += 1
-
-
-@ngjit
-def aa_stage_2_clear(aggs_and_copies, antialias_zeroes):
-    k = 0
-    # Numba access to heterogeneous tuples is only permitted using literal_unroll.
-    for agg_and_copy in literal_unroll(aggs_and_copies):
-        parallel_fill(agg_and_copy[0], antialias_zeroes[k])
-        k += 1
-
-
-@ngjit
-def aa_stage_2_copy_back(aggs_and_copies):
-    # Numba access to heterogeneous tuples is only permitted using literal_unroll.
-    for agg_and_copy in literal_unroll(aggs_and_copies):
-        agg_and_copy[0][:] = agg_and_copy[1][:]

datashader/compiler.py

@@ -3,11 +3,15 @@
 from typing import TYPE_CHECKING
 
 from toolz import unique, concat, pluck, get, memoize
+from numba import literal_unroll
 import numpy as np
 import xarray as xr
 
+from .antialias import AntialiasCombination
 from .reductions import SpecialColumn, by, category_codes, summary
-from .utils import isnull, ngjit
+from .utils import (isnull, ngjit, parallel_fill, nanmax_in_place, nanmin_in_place, nansum_in_place,
+    nanfirst_in_place, nanlast_in_place,
+)
 
 try:
     from datashader.transfer_functions._cuda_utils import cuda_mutex_lock, cuda_mutex_unlock
@@ -75,7 +79,13 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False, parti
     ``antialias_stage_2``
         If using antialiased lines this is a tuple of the ``AntialiasCombination``
         values corresponding to the aggs. If not using antialiased lines then
-        this is False.
+        this is ``False``.
+
+    ``antialias_stage_2_funcs``
+        If using antialiased lines which require a second stage combine, this
+        is a tuple of the three combine functions which are the accumulate,
+        clear and copy_back functions. If not using antialiased lines then this
+        is ``None``.
 
     ``column_names``
         Names of DataFrame columns or DataArray variables that are used by the
@@ -99,9 +109,11 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False, parti
         else:
             array_module = np
         antialias_stage_2 = antialias_stage_2(array_module)
+        antialias_stage_2_funcs = make_antialias_stage_2_functions(antialias_stage_2)
     else:
         self_intersect = False
         antialias_stage_2 = False
+        antialias_stage_2_funcs = None
 
     # List of tuples of
     # (append, base, input columns, temps, combine temps, uses cuda mutex, is_categorical)
@@ -125,7 +137,64 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False, parti
 
     column_names = [c.column for c in cols if c.column != SpecialColumn.RowIndex]
 
-    return create, info, append, combine, finalize, antialias_stage_2, column_names
+    return create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, column_names
+
+
+def _get_antialias_stage_2_combine_func(combination: AntialiasCombination, zero: float,
+                                        n_reduction: bool, categorical: bool):
+    if not n_reduction:
+        # The aggs to combine here are either 3D (ny, nx, ncat) if categorical is True or
+        # 2D (ny, nx) if categorical is False. The same combination functions can be for both
+        # as all elements are independent.
+        if combination == AntialiasCombination.MAX:
+            return nanmax_in_place
+        elif combination == AntialiasCombination.MIN:
+            return nanmin_in_place
+        elif combination == AntialiasCombination.FIRST:
+            return nanfirst_in_place
+        elif combination == AntialiasCombination.LAST:
+            return nanlast_in_place
+        else:
+            return nansum_in_place
+
+    raise NotImplementedError
+
+
+def make_antialias_stage_2_functions(antialias_stage_2):
+    aa_combinations, aa_zeroes, aa_n_reductions, aa_categorical = antialias_stage_2
+
+    # Accumulate functions.
+    funcs = [_get_antialias_stage_2_combine_func(comb, zero, n_red, cat) for comb, zero, n_red, cat
+             in zip(aa_combinations, aa_zeroes, aa_n_reductions, aa_categorical)]
+
+    namespace = {}
+    namespace["literal_unroll"] = literal_unroll
+    for func in set(funcs):
+        namespace[func.__name__] = func
+
+    lines = ["def aa_stage_2_accumulate(aggs_and_copies):"]
+    for i, func in enumerate(funcs):
+        lines.append(f"    {func.__name__}(aggs_and_copies[{i}][1], aggs_and_copies[{i}][0])")
+
+    code = "\n".join(lines)
+    exec(code, namespace)
+    aa_stage_2_accumulate = ngjit(namespace["aa_stage_2_accumulate"])
+
+    @ngjit
+    def aa_stage_2_clear(aggs_and_copies):
+        k = 0
+        # Numba access to heterogeneous tuples is only permitted using literal_unroll.
+        for agg_and_copy in literal_unroll(aggs_and_copies):
+            parallel_fill(agg_and_copy[0], aa_zeroes[k])
+            k += 1
+
+    @ngjit
+    def aa_stage_2_copy_back(aggs_and_copies):
+        # Numba access to heterogeneous tuples is only permitted using literal_unroll.
+        for agg_and_copy in literal_unroll(aggs_and_copies):
+            agg_and_copy[0][:] = agg_and_copy[1][:]
+
+    return aa_stage_2_accumulate, aa_stage_2_clear, aa_stage_2_copy_back
 
 
 def traverse_aggregation(agg):

datashader/data_libraries/dask.py

@@ -71,11 +71,12 @@ def default(glyph, df, schema, canvas, summary, *, antialias=False, cuda=False):
 
     # Compile functions
     partitioned = isinstance(df, dd.DataFrame) and df.npartitions > 1
-    create, info, append, combine, finalize, antialias_stage_2, column_names = compile_components(
-        summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=partitioned)
+    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, column_names = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=partitioned)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
     x_range = bounds[:2]
     y_range = bounds[2:]
 
@@ -209,11 +210,12 @@ def line(glyph, df, schema, canvas, summary, *, antialias=False, cuda=False):
 
     # Compile functions
     partitioned = isinstance(df, dd.DataFrame) and df.npartitions > 1
-    create, info, append, combine, finalize, antialias_stage_2, _ = compile_components(
-        summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=partitioned)
+    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=partitioned)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
     x_range = bounds[:2]
     y_range = bounds[2:]
 

datashader/data_libraries/dask_xarray.py

@@ -58,11 +58,12 @@ def dask_rectilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False,
     shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize, antialias_stage_2, _ = compile_components(
-        summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
+    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
     x_range = bounds[:2]
     y_range = bounds[2:]
 
@@ -140,11 +141,12 @@ def dask_raster(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=
     shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize, antialias_stage_2, _ = compile_components(
-        summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
+    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
     x_range = bounds[:2]
     y_range = bounds[2:]
 
@@ -234,11 +236,12 @@ def dask_curvilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False,
     shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize, antialias_stage_2, _ = compile_components(
-        summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
+    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
     x_range = bounds[:2]
     y_range = bounds[2:]
 

datashader/data_libraries/pandas.py

@@ -23,11 +23,12 @@ def pandas_pipeline(df, schema, canvas, glyph, summary, *, antialias=False):
 @glyph_dispatch.register(_GeometryLike)
 @glyph_dispatch.register(_AreaToLineLike)
 def default(glyph, source, schema, canvas, summary, *, antialias=False, cuda=False):
-    create, info, append, _, finalize, antialias_stage_2, _ = compile_components(
-        summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=False)
+    create, info, append, _, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=False)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
 
     x_range = canvas.x_range or glyph.compute_x_bounds(source)
     y_range = canvas.y_range or glyph.compute_y_bounds(source)

datashader/glyphs/area.py

@@ -98,7 +98,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -182,7 +182,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -274,7 +274,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -365,7 +365,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -472,7 +472,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -580,7 +580,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -649,7 +649,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -728,7 +728,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -797,7 +797,7 @@ def compute_bounds_dask(self, ddf):
                 self.compute_y_bounds())
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -865,7 +865,7 @@ def compute_bounds_dask(self, ddf):
                 self.compute_y_bounds())
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -950,7 +950,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -1030,7 +1030,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
+    def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(