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


class numba_dask_decorator:
    def __init__(self, dim_tracer: ak.Array):
        self.dim_tracer = dim_tracer

    def __call__(self,callable):
        self.jit_method = numba.jit(parallel=False, fastmath=False)(callable)

        def _wrapped_method(*args):
            def call_underlying(*args):
                """Calling a numba JIT version of the underlying function"""
                if ak.backend(args[0]) == "typetracer":
                    _ = [ak.typetracer.touch_data(x) for x in args]
                    return ak.Array(
                        self.dim_tracer.layout.to_typetracer(forget_length=True)
                    )
                return self.jit_method(*args, ak.ArrayBuilder()).snapshot()

            if isinstance(args[0], ak.Array):
                return call_underlying(*args)
            elif isinstance(args[0], dak.Array):
                return dak.lib.core.map_partitions(
                    call_underlying,
                    *args,
                    label="dask-numba-" + callable.__name__,
                )
        return _wrapped_method


@numba_dask_decorator(dim_tracer=ak.Array([[0.0]]))
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


@numba_dask_decorator(dim_tracer=ak.Array([[0.0]]))
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


@numba_dask_decorator(dim_tracer=ak.Array([0.0]))
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.

Update 2025.Apr.07

Now explicitly requesting the function defining to provide the expected output dimension

[lgray]

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