Вызов GSL из Numba без ctypes и cffi с кэшированием

Показываем, как вызывать GSL из Numba без ctypes/cffi: загрузка через llvmlite, ExternalFunction, упаковка gsl_function и кортежей, кэшируемые вызовы qag.

Оборачивать GSL из чистого Python и при этом сохранять кэшируемость Numba обычно мешают две вещи: вызов функций из общей библиотеки без ctypes/cffi и передача структурированного состояния в колбэки GSL без самописного C‑клея. Ниже — компактный путь, который закрывает оба вопроса. Он показывает, как напрямую вызывать GSL через llvmlite и numba.types.ExternalFunction и как упаковывать аргументы для gsl_function, понижая Numba Tuple до того же ABI, что и простой C‑структуры.

Постановка задачи

Цель — обернуть выбранные процедуры GSL, по возможности обойтись без ручного C‑кода и сохранить результат кэшируемым с njit(cache=True). Базовая схема часто стартует так и использует ctypes, чтобы подгрузить небольшой C‑шим:

// gsl_bridge.c
#include <gsl/gsl_integration.h>
#include <stdint.h>

typedef double (*gsl_cb_t)(double x, void* params);

double qag_bridge(
        gsl_cb_t cb,
        void* user_state,
        int which,
        double a,
        double b,
        double epsabs,
        double epsrel,
        int limit,
        int key) {
    gsl_integration_workspace* work = gsl_integration_workspace_alloc(limit);
    double res, err;

    gsl_function fn_proxy;
    fn_proxy.function = cb;
    fn_proxy.params = user_state;

    gsl_integration_qag(&fn_proxy, a, b, epsabs, epsrel, limit, key, work, &res, &err);
    gsl_integration_workspace_free(work);

    return res;
}

# integration.py
import ctypes as ct
from pathlib import Path
from numba import njit

def find_ext(lib_base: str) -> str:
    here = Path(__file__).parent
    patt = f"*{lib_base}.*"
    found = here.glob(patt)
    try:
        return str(next(found))
    except StopIteration:
        return None

_so_path = find_ext('gsl_integration')
_so_handle = ct.CDLL(_so_path)

qag_fn = _so_handle.qag
qag_fn.argtypes = [ct.c_void_p, ct.c_void_p,
                   ct.c_double, ct.c_double,
                   ct.c_double, ct.c_double,
                   ct.c_int, ct.c_int]
qag_fn.restype = ct.c_double

@njit
def qag_jit(cb_addr: int,
            a: float,
            b: float,
            epsabs: float = 1.49e-08,
            epsrel: float = 1.49e-08,
            limit: int = 50,
            key: int = 1,
            user_ptr: int = 0) -> float:
    return qag_fn(cb_addr, user_ptr, a, b, epsabs, epsrel, limit, key)

# setup.py
from setuptools import setup, Extension, find_packages

ext_mod = Extension(
    "numba_gsl.gsl_integration",
    sources=["numba_gsl/gsl_integration.c"],
    libraries=["gsl", "gslcblas"],
    extra_compile_args=["-O3"],
)

setup(
    name="numba_gsl",
    version="0.1.0",
    description="Numba-friendly wrappers for GSL routines",
    author="No Name",
    packages=find_packages(),
    ext_modules=[ext_mod],
    install_requires=["numba", "llvmlite"],
    python_requires=">=3.12",
)

# demo.py
import math
from numba import cfunc, types
from numba_gsl.integration import qag_jit
from scipy.integrate import quad

@cfunc(types.float64(types.float64, types.voidptr))
def sinc_cb(x, _):
    return math.sin(x) / x if x != 0.0 else 1.0

def py_sinc(x):
    return math.sin(x) / x if x != 0.0 else 1.0

addr = sinc_cb.address
res_numba = qag_jit(addr, a=1e-8, b=3.14)
res_scipy = quad(py_sinc, a=1e-8, b=3.14)[0]

print("numba_gsl quad result:", res_numba)
print("scipy     quad result:", res_scipy)

Подход срабатывает в простых случаях, но опирается на ctypes. Это блокирует njit(cache=True) и оставляет сложные структуры GSL вроде gsl_function и workspace на откуп Python‑коду или одноразовым C‑шимам, что ограничивает расширяемость.

Откуда берётся трение

Возникают две независимые проблемы. Первая — как вызывать функции из DLL/SO из Numba так, чтобы код кэшировался. Эту часть решает единоразовая загрузка библиотеки через llvmlite.binding.load_library_permanently и привязка символов через numba.types.ExternalFunction вместо ctypes или cffi. Вторая — как передать в GSL колбэк и блок параметров, который на стороне C выглядит как gsl_function. Кортеж Numba (Tuple) понижается в LLVM до того же представления, что и C‑структура. Значит, упакованный кортеж как маленькая структура с приведением к void* может играть роль полезной нагрузки для gsl_function. Делать это нужно аккуратно: компиляторы C могут по-разному раскладывать сложные структуры в ABI. Но для простых форм — например, пары указателей — раскладка обычно стабильна на разных платформах.

Решение: кэшируемая линковка и совместимое с ABI понижение кортежей

Дальнейшие шаги опираются на три идеи. Загрузить стороннюю библиотеку один раз в процесс, привязать C‑символы как ExternalFunction и мостить значения к указателям и обратно небольшими intrinsic‑функциями. Вспомогательная функция превращает Tuple в «псевдо‑указатель», чтобы GSL мог положить его в поле params структуры gsl_function.

from numba import types, typeof, njit, cfunc
from numba.extending import intrinsic
from numba.core import cgutils
import numba as nb
import numpy as np

@intrinsic
def stack_ptr_from_val(typingctx, item):
    def impl(context, builder, signature, args):
        tmp = cgutils.alloca_once_value(builder, args[0])
        return tmp
    sig = types.CPointer(typeof(item).instance_type)(typeof(item).instance_type)
    return sig, impl

@intrinsic
def deref_ptr(typingctx, arrptr):
    def impl(context, builder, signature, args):
        val = builder.load(args[0])
        return val
    sig = arrptr.dtype(types.CPointer(arrptr.dtype))
    return sig, impl

@intrinsic
def sizeof_tuple_abi(typingctx, tup_t):
    if not isinstance(tup_t, types.Tuple):
        raise TypeError("Argument must be a tuple")

    def codegen(context, builder, sig, args):
        ty = context.get_value_type(tup_t)
        base_ptr = cgutils.alloca_once(builder, ty)
        next_ptr = builder.gep(base_ptr, [context.get_constant(types.int32, 1)])
        intptr = context.get_value_type(types.intp)
        base = builder.ptrtoint(base_ptr, intptr)
        nxt = builder.ptrtoint(next_ptr, intptr)
        sz = builder.sub(nxt, base)
        return sz

    sig = types.intp(tup_t)
    return sig, codegen

@intrinsic
def tuple_as_voidptr(typingctx, tup_t):
    if not isinstance(tup_t, types.Tuple):
        raise TypeError("Argument must be a tuple")

    sig = types.voidptr(tup_t)

    def codegen(context, builder, signature, args):
        (val,) = args
        slot = cgutils.alloca_once_value(builder, val)
        void_ty = context.get_value_type(types.voidptr)
        casted = builder.bitcast(slot, void_ty)
        return casted

    return sig, codegen

def make_cfunc_and_sig(fn, arg_pack, cache=False):
    if not callable(fn) and not isinstance(fn, nb.core.registry.CPUDispatcher):
        raise TypeError("fn must be callable")
    if not isinstance(arg_pack, tuple):
        raise TypeError("arg_pack must be a tuple")
    if isinstance(fn, nb.core.registry.CPUDispatcher):
        fn = fn.py_func

    nb_inner = nb.njit(types.double(types.double, nb.typeof(arg_pack)), inline="always", cache=cache)(fn)
    sig = types.double(types.double, types.CPointer(nb.typeof(arg_pack)))

    @cfunc(sig, error_model="numpy", cache=True)
    def cfun(x, pack_ptr):
        return nb_inner(x, pack_ptr[0])

    return cfun, sig

def make_cfunc_with_blob(fn, arg_pack, cache=False):
    if not callable(fn) and not isinstance(fn, nb.core.registry.CPUDispatcher):
        raise TypeError("fn must be callable")
    if not isinstance(arg_pack, tuple):
        raise TypeError("arg_pack must be a tuple")
    if isinstance(fn, nb.core.registry.CPUDispatcher):
        fn = fn.py_func

    @njit(cache=cache)
    def lower_tuple_blob(tup):
        size = sizeof_tuple_abi(tup)
        blob_ptr = tuple_as_voidptr(tup)
        view = nb.carray(blob_ptr, (size,), dtype=np.uint8)
        copy = np.copy(view)
        return copy

    nb_inner = nb.njit(types.double(types.double, nb.typeof(arg_pack)), inline="always", cache=cache)(fn)

    @cfunc(types.double(types.double, types.CPointer(nb.typeof(arg_pack))), error_model="numpy", cache=cache)
    def cfun(x, pack_ptr):
        return nb_inner(x, pack_ptr[0])

    return cfun, lower_tuple_blob

@njit(inline="always")
def pack_gsl_fun(cb_addr, arg_pack):
    arg_ptr = tuple_as_voidptr(arg_pack)
    return tuple_as_voidptr((cb_addr, arg_ptr))

С этими intrinsic‑функциями можно напрямую привязываться к GSL через load_library_permanently и ExternalFunction. Это позволяет кэшировать код и полностью обойтись без ctypes.

from llvmlite import binding
from numba import types
import numba as nb
import gsl_intrinsics as gi

binding.load_library_permanently(r'path_to_gsl.dll')

f64   = types.double
f64_p = types.CPointer(f64)
i32   = types.intc
sz_t  = types.size_t
vptr  = types.voidptr

name_qag = 'gsl_integration_qag'
sig_qag = i32(
    vptr,   # gsl_function*
    f64,    # a
    f64,    # b
    f64,    # epsabs
    f64,    # epsrel
    sz_t,   # limit
    i32,    # key
    vptr,   # workspace*
    f64_p,  # result*
    f64_p   # abserr*
)
qag_ext = types.ExternalFunction(name_qag, sig_qag)

name_alloc = 'gsl_integration_workspace_alloc'
sig_alloc = vptr(sz_t,)
ws_alloc = types.ExternalFunction(name_alloc, sig_alloc)

name_free = 'gsl_integration_workspace_free'
sig_free = types.none(vptr,)
ws_free = types.ExternalFunction(name_free, sig_free)

@nb.njit(cache=True)
def call_gsl_qag(cb_addr, arg_pack, a, b, epsabs, epsrel, limit, key):
    params_ptr = gi.tuple_as_voidptr(arg_pack)
    gsl_fn = gi.tuple_as_voidptr((cb_addr, params_ptr))

    work = ws_alloc(limit)

    res_ptr = gi.stack_ptr_from_val(nb.double(0))
    err_ptr = gi.stack_ptr_from_val(nb.double(0))

    _ret = qag_ext(
        gsl_fn, a, b,
        epsabs, epsrel, limit,
        key, work,
        res_ptr, err_ptr,
    )

    ws_free(work)

    return gi.deref_ptr(res_ptr), gi.deref_ptr(err_ptr)

Использование повторяет версию на ctypes, но целиком идёт через внешний вызов Numba и потому кэшируется. Для наглядности один вариант строит cfunc из функции njit и конкретного кортежа аргументов; другой — задаёт сигнатуру cfunc вручную.

import numba as nb
from numba import types
import numpy as np
import gsl_intrinsics as gi
import math

arg_tuple = (5.0, np.ones((3, 3)), 8.0)

@nb.njit()
def sinc_plus_bias(x, pack):
    return (math.sin(x) / x + pack[2]) if x != 0.0 else 1.0

cfun_auto, sig_auto = make_cfunc_and_sig(sinc_plus_bias, arg_tuple, cache=True)

inner = (types.float64, types.Array(types.float64, 2, 'C'), types.float64)
@nb.cfunc(types.double(types.double, types.CPointer(types.Tuple(inner))), cache=True)
def cfun_manual(x, p):
    return sinc_plus_bias(x, p[0])

a = 1e-8
b = 3.14
epsabs = 1.49e-08
epsrel = 1.49e-08
limit = 50
key = 1

print(call_gsl_qag(cfun_manual.address, arg_tuple, a, b, epsabs, epsrel, limit, key))

print("Cache hits call_gsl_qag:", call_gsl_qag.stats.cache_hits)
print("Cache misses call_gsl_qag:", call_gsl_qag.stats.cache_misses)

print("")
print("Cache hits cfun_manual:", cfun_manual.cache_hits)

Почему работает трюк с кортежами

Numba понижает Tuple до LLVM‑структуры, и простые C‑структуры также становятся LLVM‑структурами. Когда раскладка проста — например, пара указателей, — выбранный ABI у распространённых C‑компиляторов обычно совпадает с тем, что нужно на уровне LLVM. Приведение Tuple к void* даёт указатель, который GSL может хранить в gsl_function.params. Сам указатель на функцию передаётся рядом и упаковывается тем же способом. Прямая привязка через ExternalFunction позволяет Numba рассматривать вызов как внешний скомпилированный и кэшировать итоговый машинный код.

Зачем это знать

Это открывает возможность строить кэшируемые пайплайны на njit, которые обращаются к вендорным библиотекам без накладных расходов ctypes/cffi и без ручной поддержки клея под каждую C‑структуру. Решение остаётся расширяемым: настроив упаковку кортежей и маленькие вспомогательные функции для указателей, добавлять новые вызовы GSL можно по той же схеме. Важное ограничение: хотя простые формы «кортеж‑как‑структура» обычно надёжны, сложные C‑структуры зависят от компилятора и платформы, так что выход за пределы тривиальных раскладок стоит делать осознанно.

Итоги

Если вам нужна кэшируемость Numba, не вызывайте через ctypes или cffi. Загрузите разделяемую библиотеку один раз с llvmlite.binding.load_library_permanently и привязывайте функции через numba.types.ExternalFunction. Чтобы передать состояние в колбэки GSL, понижайте Numba Tuple до структуроподобного блока и приводите к void*. Держите форму структуры простой, чтобы не нарваться на проблемы ABI. Если хочется более традиционного пути, низкоуровневые интерфейсы можно генерировать и через Cython — для случаев с насыщенными C‑структурами это уменьшит объём Python‑обвязки и будет практичным выбором.