Generate custom kernels with numba

[skrah]

There appears to be some support for generating custom kernels with Numba that use the C calling convention and can therefore be directly inserted into the gumath function table.

This is an example for a Strided kernel:

import gumath as gm
from xnd import xnd
from numba import cfunc, carray
from numba.types import CPointer, float64, int64, int32, void, intp, char
import sys
import numpy as np


@cfunc(int32(CPointer(CPointer(float64)), CPointer(intp), CPointer(intp), CPointer(void)), nopython=True)
def absolute__d_d(args, dimensions, steps, data):
    src = args[0]
    dest = args[1]
    N = dimensions[0]
    step = steps[0] // 8
    i = 0
    for k in range(N):
        dest[k] = abs(src[i])
        i += step
    return 0

# Get function pointer and insert kernel into the lookup table.
ptr = absolute__d_d.address
gm.abs = gm.unsafe_add_numpy_kernel(name="abs", sig="... * float64 -> ... * float64", ptr=ptr)


x = xnd([-1.0, -2e122, 3.0])
y = gm.abs(x)
print(y)

x = xnd([-1.0, -2e122, 3.0])
y = gm.abs(x[::-1])
print(y)

However, this example works because all arguments are float64. Otherwise the first argument should be CPointer(CPointer(char)) and one would need to cast to e.g. s = CPointer(int8(src))) inside the function.

The latter does not appear to be supported -- casting to primitive numpy types works.

[skrah]

Using this method for Xnd kernels seems to fail:

import gumath as gm
from xnd import xnd
import numpy as np
from numba import cfunc, jit, carray, numpy_support
from numba.types import CPointer, float64, int64, int32, void, intp, char
import sys


# Simpler way to create structs?
np_xnd_bitmap = np.dtype([('data', 'p'),
                          ('size', np.float64),
                          ('next', 'p')], align=True)

np_xnd = np.dtype([('bitmap', np_xnd_bitmap),
                   ('index', np.int64),
                   ('type', 'p'),
                   ('ptr', 'p')], align=True)

nb_xnd_bitmap = numpy_support.from_dtype(np_xnd_bitmap)
nb_xnd = numpy_support.from_dtype(np_xnd)


@cfunc(int32(CPointer(nb_xnd), CPointer(void)), nopython=True)
def absolute__d_d(stack, ctx):
    src = stack[0].ptr # Numba.errors.LoweringError: Failed at nopython (nopython mode backend)
    return 0

[skrah]

In case anyone wonders why one would use gumath at all here:

1. If these examples work, gumath will do all the work of type-checking and dispatching, so large parts of internal Numba code could be simplified.

2. Xnd kernels are more general than numpy kernels, so additional features like ragged arrays could be used.