2025, Nov 16 15:03
Вызов 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‑обвязки и будет практичным выбором.