cache numba's functions

[mikilterribile]

Hi everyone,

I have been using Datashader with Matplotlib to create dynamic visualizations of large vectors (shapefiles). The approach works fine, but each time I start the program, Numba's functions for creating the canvas and the shaded image are compiled by Numba, resulting in a 6-7 second overhead.

I wonder if it's possible to enable caching (for instance, setting use_cache=True) in the @njit decorator for the function, and whether there is a general solution to overcome this problem.

Thank you in advance for your attention.

Best regards,
Michele Zucchelli.

[jbednar]

Yes, this is an issue. Definitely needs addressing!

[mikilterribile]

ok thank you for your answer

[kevinjohncutler]

This would be really useful!

[jbednar]

Yep! Would be great if someone could try turning it on and document the results here. Last time I tried that (probably years ago by now) it didn't work, but Numba's been updated a lot since then, and it may well be in a better state!

[kevinjohncutler]

I've been working on it for a few hours and it is slow going. Simply using cache=True for the global decorator @ngjit in utils.py does not work. A lot of functions that I need to cache are dynamically created, and for me these are the ones creating a 7-8s overhead:

 time(s) count  function (source file)
--------------------------------------------------------------------------------
   7.444     7  datashader.glyphs.line.extend_cpu  [<extend_cpu>]
   6.343     7  datashader.glyphs.line.perform_extend_line  [<perform_extend_line>]
   5.325     7  datashader.glyphs.line.draw_segment  [<draw_segment>]
   3.972     7  datashader.glyphs.line._full_antialias  [<_full_antialias>]
   0.266     1  datashader.glyphs.line._liang_barsky  [/home/kcutler/DataDrive/datashader/datashader/glyphs/line.py]
   0.162     3  datashader.glyphs.line._clipt  [/home/kcutler/DataDrive/datashader/datashader/glyphs/line.py]

From what I gather, the top-level decorated functions have a defined code.co_filename that allows numba to find the source code and check the hash. Each inner definition of extend_cpu, for example, does not follow this structure and has no locator for numba to check. I'm trying to refactor _full_antialias now to give it a definite locator, but no luck yet.

Related: numba/numba#3501

[kevinjohncutler]

An alternative that might work for some users might be to simply compile the relevant dynamic functions in the background on package import, such that it's done by the time datashader is invoked. In my case, there are multiple shapes of data that are rendered in a single line plot, and each requires its own specialized function to compile in datashader (also line width and dtype). However, even with multithreading, I found that this does not complete in short enough timespan to be useful.

[jbednar]

That sounds like what I found last time, a couple of years back. At that time, what I was wondering was whether we could avoid the problem by using newer features in Numba so that we could avoid having to generate specialized code dynamically. The hope was that Numba made some good progress since the original 2016 Datashader code. But (a) I don't know if Numba addressed the underlying constraint we had, and (b) making use of any improvements to Numba would require rewriting some of the hairiest internal bits of Datashader.

[kevinjohncutler]

Seems like there have not been any relevant improvements to Numba on that front. At least two abandoned PRs tried to get at this issue. The more successful workaround I linked to above simply exports the dynamically generated code to real on-disk files that numba can locate. I also tried this approach to replace the exec calls in compile.py but that doesn't seem to touch any of the functions in line.py.

It might be worth seeing if the former approach works. A little refactoring happened on that repo since that original post was made. Relevant function definitions are now here.

[kevinjohncutler]

Might as well loop in @DannyWeitekamp to see if his approach will likely work here and get a sense of how hard it will be. I adapted one function in line.py to use the cre_caching functions, but that was a much simpler case where njit(cache=True) would also have worked just as well. The nested functions in lines.py are another beast entirely. For example:

  
def _build_extend_line_axis1_x_constant(draw_segment, expand_aggs_and_cols,
                                        antialias_stage_2_funcs, swap_dims: bool = False):
    if antialias_stage_2_funcs is not None:
        aa_stage_2_accumulate, aa_stage_2_clear, aa_stage_2_copy_back = antialias_stage_2_funcs
    use_2_stage_agg = antialias_stage_2_funcs is not None
    
    @ngjit
    @expand_aggs_and_cols
    def perform_extend_line(
            i, j, sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys, buffer, *aggs_and_cols
    ):
        x0 = xs[j]
        x1 = xs[j + 1]
        if swap_dims:
            y0 = ys[j, i]
            y1 = ys[j + 1, i]
            segment_start = (j == 0) or isnull(xs[j - 1]) or isnull(ys[j - 1, i])
            segment_end = (j == len(xs)-2) or isnull(xs[j+2]) or isnull(ys[j+2, i])
        else:
            y0 = ys[i, j]
            y1 = ys[i, j + 1]
            segment_start = (j == 0) or isnull(xs[j - 1]) or isnull(ys[i, j - 1])
            segment_end = (j == len(xs)-2) or isnull(xs[j+2]) or isnull(ys[i, j+2])

        if segment_start or use_2_stage_agg:
            xm = 0.0
            ym = 0.0
        else:
            xm = xs[j-1]
            ym = ys[j-1, i] if swap_dims else ys[i, j-1]

        draw_segment(i, sx, tx, sy, ty, xmin, xmax, ymin, ymax,
                     segment_start, segment_end, x0, x1, y0, y1,
                     xm, ym, buffer, *aggs_and_cols)

    @ngjit
    @expand_aggs_and_cols
    def extend_cpu(sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys, antialias_stage_2,
                   *aggs_and_cols):
        antialias = antialias_stage_2 is not None
        buffer = np.empty(8) if antialias else None
        ncols, nrows = ys.shape if swap_dims else ys.shape[::-1]
        for i in range(nrows):
            for j in range(ncols - 1):
                perform_extend_line(
                    i, j, sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys, buffer, *aggs_and_cols
                )

    def extend_cpu_antialias_2agg(sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys,
                                  antialias_stage_2, *aggs_and_cols):
        n_aggs = len(antialias_stage_2[0])
        aggs_and_accums = tuple((agg, agg.copy()) for agg in aggs_and_cols[:n_aggs])
        cpu_antialias_2agg_impl(sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys,
                                antialias_stage_2, aggs_and_accums, *aggs_and_cols)

    @ngjit
    @expand_aggs_and_cols
    def cpu_antialias_2agg_impl(sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys,
                                antialias_stage_2, aggs_and_accums, *aggs_and_cols):
        antialias = antialias_stage_2 is not None
        buffer = np.empty(8) if antialias else None

        ncols = ys.shape[1]
        for i in range(ys.shape[0]):
            for j in range(ncols - 1):
                perform_extend_line(
                    i, j, sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys,
                    buffer, *aggs_and_cols
                )

            if ys.shape[0] == 1:
                return

            aa_stage_2_accumulate(aggs_and_accums, i==0)

            if i < ys.shape[0] - 1:
                aa_stage_2_clear(aggs_and_accums)

        aa_stage_2_copy_back(aggs_and_accums)

    @cuda.jit
    @expand_aggs_and_cols
    def extend_cuda(sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys, antialias_stage_2,
                    *aggs_and_cols):
        antialias = antialias_stage_2 is not None
        buffer = cuda.local.array(8, nb_types.float64) if antialias else None
        i, j = cuda.grid(2)
        ncols, nrows = ys.shape if swap_dims else ys.shape[::-1]
        if i < nrows and j < ncols - 1:
            perform_extend_line(
                i, j, sx, tx, sy, ty, xmin, xmax, ymin, ymax, xs, ys, buffer, *aggs_and_cols
            )

    if use_2_stage_agg:
        return extend_cpu_antialias_2agg, extend_cuda
    else:
        return extend_cpu, extend_cuda

[philippjfr]

Not sure this is immediately helpful here but we've long hoped we could back out of a lot of the dynamic code generation. A lot of it was necessary in earlier versions of Numba because performance degraded massively if the code wasn't structured just right, e.g. of the top of my head explicitly listing out all the function arguments was much more efficient than argument unpacking. The hope was that modern Numba is smarter and we can simplify a lot of that, but we haven't had the bandwidth to investigate if that's true or at all feasible.