jan/klamath-rs
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

CFFI + maturin approach

Browse source
This commit is contained in:
jan 2024-12-21 15:42:50 -08:00
commit bac31e2ce6
3 changed files with 92 additions and 130 deletions
Download patch file Download diff file Expand all files Collapse all files

75
klamath_rs_ext/basic.py
View file

@ -6,62 +6,47 @@ from itertools import chain
import numpy
from numpy.typing import NDArray
so_path = Path(__file__).resolve().parent / 'libklamath_rs_ext.so'
clib = ctypes.CDLL(so_path)
from .klamath_rs_ext import lib, ffi
CONV_TABLE_i16 = {
numpy.float64: clib.f64_to_i16,
numpy.float32: clib.f32_to_i16,
numpy.int64: clib.i64_to_i16,
numpy.int32: clib.i32_to_i16,
numpy.int16: clib.i16_to_i16,
numpy.uint64: clib.u64_to_i16,
numpy.uint32: clib.u32_to_i16,
numpy.uint16: clib.u16_to_i16,
numpy.float64: lib.f64_to_i16,
numpy.float32: lib.f32_to_i16,
numpy.int64: lib.i64_to_i16,
numpy.int32: lib.i32_to_i16,
numpy.int16: lib.i16_to_i16,
numpy.uint64: lib.u64_to_i16,
numpy.uint32: lib.u32_to_i16,
numpy.uint16: lib.u16_to_i16,
}
CONV_TABLE_i32 = {
numpy.float64: clib.f64_to_i32,
numpy.float32: clib.f32_to_i32,
numpy.int64: clib.i64_to_i32,
numpy.int32: clib.i32_to_i32,
numpy.uint64: clib.u64_to_i32,
numpy.uint32: clib.u32_to_i32,
numpy.float64: lib.f64_to_i32,
numpy.float32: lib.f32_to_i32,
numpy.int64: lib.i64_to_i32,
numpy.int32: lib.i32_to_i32,
numpy.uint64: lib.u64_to_i32,
numpy.uint32: lib.u32_to_i32,
}
clib.f64_to_i16.restype = ctypes.c_double
clib.f32_to_i16.restype = ctypes.c_float
clib.i64_to_i16.restype = ctypes.c_int64
clib.i32_to_i16.restype = ctypes.c_int32
clib.i16_to_i16.restype = ctypes.c_int16
clib.u64_to_i16.restype = ctypes.c_uint64
clib.u32_to_i16.restype = ctypes.c_uint32
clib.u16_to_i16.restype = ctypes.c_uint16
clib.f64_to_i32.restype = ctypes.c_double
clib.f32_to_i32.restype = ctypes.c_float
clib.i64_to_i32.restype = ctypes.c_int64
clib.i32_to_i32.restype = ctypes.c_int32
clib.u64_to_i32.restype = ctypes.c_uint64
clib.u32_to_i32.restype = ctypes.c_uint32
for fn in chain(CONV_TABLE_i16.values(), CONV_TABLE_i32.values()):
fn.argtypes = [ctypes.POINTER(fn.restype), ctypes.c_size_t]
def pack_int2(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
arr = numpy.asarray(data)
if arr.dtype in CONV_TABLE_i16.keys():
for dtype in CONV_TABLE_i16.keys():
if arr.dtype != dtype:
continue
arr = numpy.require(arr, requirements=('C_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', 'OWNDATA'))
if arr is data:
arr = numpy.array(arr, copy=True)
fn = CONV_TABLE_i16[arr.dtype]
result = fn(arr.ctypes.data_as(fn.argtypes[0]), arr.size)
fn = CONV_TABLE_i16[dtype]
result = fn(ffi.from_buffer(arr), arr.size)
if result != 0:
raise ValueError(f'Invalid value for conversion to Int2: {result}')
i2arr = arr.view('>i2')[::arr.itemsize // 2]
return i2arr.tobytes()
@ -78,13 +63,19 @@ def pack_int2(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
def pack_int4(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
arr = numpy.asarray(data)
if arr.dtype in CONV_TABLE_i32.keys():
for dtype in CONV_TABLE_i32.keys():
if arr.dtype != dtype:
continue
arr = numpy.require(arr, requirements=('C_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', 'OWNDATA'))
if arr is data:
arr = numpy.array(arr, copy=True)
fn = CONV_TABLE_i32[arr.dtype]
result = fn(arr.ctypes.data_as(fn.argtypes[0]), arr.size)
fn = CONV_TABLE_i32[dtype]
result = fn(ffi.from_buffer(arr), arr.size)
if result != 0:
raise ValueError(f'Invalid value for conversion to Int4: {result}')
i4arr = arr.view('>i4')[::arr.itemsize // 4]
return i4arr.tobytes()