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

ctypes approach

Browse source
This commit is contained in:
jan 2024-12-21 13:56:51 -08:00
commit 320958d888
5 changed files with 118 additions and 95 deletions
Download patch file Download diff file Expand all files Collapse all files

2
klamath_rs_ext/__init__.py
View file

@ -2,5 +2,5 @@ from .basic import pack_int2 as pack_int2
from .basic import pack_int4 as pack_int4
__version__ = 0.1
__version__ = 0.2

71
klamath_rs_ext/basic.py
View file

@ -1,24 +1,68 @@
from collections.abc import Sequence
import ctypes
from pathlib import Path
from itertools import chain
import numpy
from numpy.typing import NDArray
from .klamath_rs_ext import arr_to_int2, arr_to_int4
so_path = Path(__file__).resolve().parent / 'libklamath_rs_ext.so'
clib = ctypes.CDLL(so_path)
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,
}
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,
}
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 (
numpy.float64, numpy.float32,
numpy.int64, numpy.uint64,
numpy.int32, numpy.uint32,
numpy.int16, numpy.uint16,
):
if arr.dtype in CONV_TABLE_i16.keys():
arr = numpy.require(arr, requirements=('C_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', 'OWNDATA'))
if arr is data:
arr = numpy.array(arr, copy=True)
arr_to_int2(arr)
fn = CONV_TABLE_i16[arr.dtype]
result = fn(arr.ctypes.data_as(fn.argtypes[0]), arr.size)
i2arr = arr.view('>i2')[::arr.itemsize // 2]
return i2arr.tobytes()
@ -34,15 +78,14 @@ 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 (
numpy.float64, numpy.float32,
numpy.int64, numpy.uint64,
numpy.int32, numpy.uint32,
):
if arr.dtype in CONV_TABLE_i32.keys():
arr = numpy.require(arr, requirements=('C_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', 'OWNDATA'))
if arr is data:
arr = numpy.array(arr, copy=True)
arr_to_int4(arr)
fn = CONV_TABLE_i32[arr.dtype]
result = fn(arr.ctypes.data_as(fn.argtypes[0]), arr.size)
i4arr = arr.view('>i4')[::arr.itemsize // 4]
return i4arr.tobytes()