Glyph rendering optimization using variable length argument expansion

[jonmmease]

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 the expand_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...

@expand_varargs(2)
def example_fn(a, b, *args):
    print(a, b)
    other_fn(a, b, *args)

would transform the function AST into a function equivalent to

def example_fn(a, b, _0, _1):
    print(a, b)
    other_fn(a, b, _0, _1)

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 inside args.

@expand_varargs(2)
def example_fn(a, b, *args):
    print(a, b, args[0])
    other_fn(a, b, *args)

⬆️ 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 the Glyph baseclass. This method inputs the append function that will perform the reduction operation, and returns an expand_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, and area 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:

• Before: https://anaconda.org/jonmmease/expand_varargs_before/notebook
• After: https://anaconda.org/jonmmease/expand_varargs_after/notebook

I tested points, line, and area glyphs using the count, sum, mean, and std 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

[jonmmease]

Update: now the Glyph.expand_aggs_and_cols method will return an identity decorator if the NUMBA_DISABLE_JIT environnmet variable is set.

This environment variable is used by numba to disable JIT compilation. With this change, functions that are wrapped with both @jit and @expand_varargs will remain unchanged when NUMBA_DISABLE_JIT is set. This makes it possible to debug these functions with a graphical debugger.

[jonmmease]

@philippjfr this is ready to go apart from that one Travis test flaking out again. Could you restart it?

Update: Never mind, looks like I can restart myself.

[jonmmease]

all green. merging