Running dask with custom numba acceleration

[yimuchen]

Created a custom decorator to allow for numba-accelerated array builder to allow users to provide computations that may be difficult to implement with awkward calls. Currently, the implementation requires that:

• The user-defined function takes ak.ArrayBuilder as a final argument
• The user-defined function must take awkward arrays with concrete types as the main methods (required for numba acceleration)
• In addition to the main user-defined computation, the user must also provide a method that return an array of the expected output dimension.

import awkward as ak
import dask_awkward as dak
import numba
import numpy as np


def numba_dask_decorator(dim_trace):
    def _decorator_wrap(callable):
        jit_callable = numba.jit(parallel=False, fastmath=False)(callable)

        def _wrapped_method(*args):
            def call_underlying(*args):
                """Calling a numba JIT version of the underlying function"""
                try:
                    return jit_callable(*args, ak.ArrayBuilder()).snapshot()
                except:
                    return dim_trace(*args)

            if isinstance(args[0], ak.Array):
                return call_underlying(*args)
            elif isinstance(args[0], dak.Array):
                args_typetrace = [
                    ak.typetracer.length_one_if_typetracer(x._meta) for x in args
                ]
                meta = call_underlying(*args_typetrace)
                return dak.lib.core.map_partitions(
                    call_underlying,
                    *args,
                    label="dask-numba-" + callable.__name__,
                    meta=ak.Array(
                        meta.layout.to_typetracer(forget_length=True),
                    ),
                )

            else:
                return

        return _wrapped_method

    return _decorator_wrap


def identical(arr):
    return ak.ones_like(arr)


@numba_dask_decorator(dim_trace=identical)
def cumsum_2d(arr, builder: ak.ArrayBuilder):
    for _a in arr:
        builder.begin_list()
        entry = 0
        for __a in _a:
            entry = entry + __a
            builder.real(entry)
        builder.end_list()
    return builder


def two_to_one(a, b):
    return a + b


@numba_dask_decorator(dim_trace=two_to_one)
def my_add(a, b, builder: ak.ArrayBuilder):
    for _a, _b in zip(a, b):
        builder.begin_list()
        for __a, __b in zip(_a, _b):
            builder.real((__a + __b) / 10)
        builder.end_list()
    return builder


def reduce_inner(a):
    return ak.sum(a, axis=-1)


@numba_dask_decorator(dim_trace=reduce_inner)
def my_reduce(a, builder: ak.ArrayBuilder):
    for _a in a:
        sum = 0
        for i, __a in enumerate(_a):
            sum = __a + (i * 10)
        builder.real(sum)
    return builder


n_entries = np.arange(0, 10)
arr = ak.zip(
    {
        "x": ak.unflatten(np.ones(np.sum(n_entries)), n_entries),
        "y": ak.unflatten(np.zeros(np.sum(n_entries)), n_entries),
    }
)

ak_ans = cumsum_2d(arr.x)
print(ak_ans)
ak_ans = my_add(arr.x, arr.y)
print(ak_ans)
ak_ans = my_reduce(arr.x)
print(ak_ans)

d_array = dak.from_awkward(arr, npartitions=3)
dak_ans = cumsum_2d(d_array.x)
print(dak_ans, dak_ans.compute(), dak.necessary_columns(dak_ans))
dak_ans = my_add(d_array.x, d_array.y)
print(dak_ans, dak_ans.compute(), dak.necessary_columns(dak_ans))
dak_ans = my_reduce(d_array.x)
print(dak_ans, dak_ans.compute(), dak.necessary_columns(dak_ans))

Potential improvements

• For some reason, the columns report for cumsum_2d is showing arr.y, though the final compute answer is correct?? I'm not sure what is going on here.
• Better way of detecting when type tracing arrays are being evaluated.
• What to do about dimension expansion when amount of expansion part of the computation?

[lgray]

Just to have all this together for when awkward-cuda has all of its interoperability bits working: NVIDIA/numba-cuda#37