Glyph rendering optimization using variable length argument expansion #780
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Background
This PR is motivated by the performance characteristics discovered while developing the quadmesh glyph in #779. In particular, see the discussion of fixed vs variable length arguments in #779 (comment).
The key insight here is that during glyph rendering, the numba optimization of the rendering functions can result in substantially faster code if the input functions contain a fixed number of arguments rather than a variable length argument (e.g. *args).
In the glyph rendering code, this variable length argument is usually called
*aggs_and_cols
and it contains a list of the aggregate arrays that are being populated and the columns that are used as input to the reduction calculations. The length of this argument varies depending on the chosen reduction operation, but the length will remain the same for every render call for a given operation.Implementation
expand_varargs
This PR adds a
datashader.macros
module that provides theexpand_varargs
decorator builder. This decorator builder inputs the desired number of arguments to expand to, and returns a function decorator. This function decorator will transform the AST of the wrapped function to replace variable length arguments with a fixed number of arguments/variables.This only makes sense for cases where the only thing the function does with the variable length argument is to pass it along to other functions in star-form. For example, calling...
would transform the function AST into a function equivalent to
If the variable length argument is used in any other context then an error is raised. For example, an error will be raised if the
example_fn
looks insideargs
.⬆️ would raise a
ValueError
.The whole point of doing this is to transform the input function before it is passed to numba's jit compilation decorator.
Glyph.expand_aggs_and_cols
A new
expand_aggs_and_cols
method has been added to theGlyph
baseclass. This method inputs theappend
function that will perform the reduction operation, and returns anexpand_varargs
decorator configured to expand the*aggs_and_cols
argument to the correct number of fixed arguments.Glyph updates
The new decorator has been added to all applicable rendering functions in the
points
,line
, andarea
glyphs. I think Trimesh could also benefit from this technique, but it doesn't quite follow the same pattern of passing around the*aggs_and_cols
argument, so that will take some refactoring.raster
doesn't currently use the same glyph/aggregation framework.Performance results
Here are some benchmark comparisons of this branch with
master
.Notebooks:
I tested
points
,line
, andarea
glyphs using thecount
,sum
,mean
, andstd
reduction operations. Here is a plot of the results:Note that the y-axis is time in seconds.
Notice how the improvements grow more significant as the reduction operation get more complex. This corresponds to the
*aggs_and_cols
argument getting longer.The speedups for the coming quadmesh glyph should be even more significant than the gains for the area glyph above.
@jbednar @philippjfr