From bac31e2ce634af549f876a1414d2aa40251209f9 Mon Sep 17 00:00:00 2001 From: jan Date: Sat, 21 Dec 2024 15:42:50 -0800 Subject: [PATCH] CFFI + maturin approach --- klamath_rs_ext/basic.py | 75 +++++++++---------- pyproject.toml | 4 + src/lib.rs | 157 ++++++++++++++++------------------------ 3 files changed, 99 insertions(+), 137 deletions(-) diff --git a/klamath_rs_ext/basic.py b/klamath_rs_ext/basic.py index 838b42f..4540fe6 100644 --- a/klamath_rs_ext/basic.py +++ b/klamath_rs_ext/basic.py @@ -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() diff --git a/pyproject.toml b/pyproject.toml index 97d888c..06aa1af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,3 +31,7 @@ dynamic = ["version"] dependencies = [ "cffi", ] + +[tool.maturin] +bindings = "cffi" + diff --git a/src/lib.rs b/src/lib.rs index 84e8551..cf8365c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,21 +7,21 @@ pub mod elements; pub mod library; -use rust_util::ToInt2BE; -use rust_util::ToInt4BE; +use byteorder::{ByteOrder, BigEndian}; +use std::mem::size_of; -macro_rules! mkfun { - ( $fname:ident, $tt:ty, $elfn:ident ) => { - #[no_mangle] - pub extern "C" fn $fname(arr: *mut $tt, size: usize) -> $tt { - let sl = unsafe { std::slice::from_raw_parts_mut(arr, size) }; +macro_rules! impl_i16be { + ( $tt:ty, $arr:ident, $size:ident ) => { + { + let sl = unsafe { std::slice::from_raw_parts_mut($arr, $size) }; for xx in sl.iter_mut() { - let res = <$tt>::$elfn(*xx); - match res { - Err(cc) => return cc, - Ok(cc) => { *xx = cc; }, - } + if *xx < i16::MIN as $tt { return *xx } + if *xx > i16::MAX as $tt { return *xx } + + let mut buf = [0; size_of::<$tt>()]; + BigEndian::write_i16(&mut buf, *xx as i16); + *xx = <$tt>::from_le_bytes(buf); } 0 as $tt } @@ -29,91 +29,58 @@ macro_rules! mkfun { } -mkfun!(f64_to_i16, f64, convert_to_i2be); -mkfun!(f32_to_i16, f32, convert_to_i2be); -mkfun!(i64_to_i16, i64, convert_to_i2be); -mkfun!(u64_to_i16, u64, convert_to_i2be); -mkfun!(i32_to_i16, i32, convert_to_i2be); -mkfun!(u32_to_i16, u32, convert_to_i2be); -mkfun!(i16_to_i16, i16, convert_to_i2be); -mkfun!(u16_to_i16, u16, convert_to_i2be); +macro_rules! impl_i32be { + ( $tt:ty, $arr:ident, $size:ident ) => { + { + let sl = unsafe { std::slice::from_raw_parts_mut($arr, $size) }; + for xx in sl.iter_mut() { + if *xx < i32::MIN as $tt { return *xx } + if *xx > i32::MAX as $tt { return *xx } -mkfun!(f64_to_i32, f64, convert_to_i4be); -mkfun!(f32_to_i32, f32, convert_to_i4be); -mkfun!(i64_to_i32, i64, convert_to_i4be); -mkfun!(u64_to_i32, u64, convert_to_i4be); -mkfun!(i32_to_i32, i32, convert_to_i4be); -mkfun!(u32_to_i32, u32, convert_to_i4be); - - -mod rust_util { - use byteorder::{ByteOrder, BigEndian}; - use std::mem::size_of; - - pub trait ToInt2BE { - fn convert_to_i2be(ii: Self) -> Result where Self: Sized; - } - - pub trait ToInt4BE { - fn convert_to_i4be(ii: Self) -> Result where Self: Sized; - } - - macro_rules! impl_i2be { - ( $tt:ty ) => { - impl ToInt2BE for $tt { - fn convert_to_i2be(ii: $tt) -> Result<$tt, $tt> { - if ii < i16::MIN as $tt { return Err(ii); } - if ii > i16::MAX as $tt { return Err(ii); } - - let mut buf = [0; size_of::<$tt>()]; - BigEndian::write_i16(&mut buf, ii as i16); - Ok(<$tt>::from_le_bytes(buf)) - } + let mut buf = [0; size_of::<$tt>()]; + BigEndian::write_i32(&mut buf, *xx as i32); + *xx = <$tt>::from_le_bytes(buf); } + 0 as $tt } } - - macro_rules! impl_i4be { - ( $tt:ty ) => { - impl ToInt4BE for $tt { - fn convert_to_i4be(ii: $tt) -> Result<$tt, $tt> { - if ii < i32::MIN as $tt { return Err(ii); } - if ii > i32::MAX as $tt { return Err(ii); } - - let mut buf = [0; size_of::<$tt>()]; - BigEndian::write_i32(&mut buf, ii as i32); - Ok(<$tt>::from_le_bytes(buf)) - } - } - } - } - - impl_i2be!(f64); - impl_i4be!(f64); - - impl_i2be!(f32); - impl_i4be!(f32); - - impl_i2be!(i64); - impl_i4be!(i64); - impl_i2be!(u64); - impl_i4be!(u64); - - impl_i2be!(i32); - impl_i4be!(i32); - impl_i2be!(u32); - impl_i4be!(u32); - - impl_i2be!(i16); - impl_i2be!(u16); - - // Does not fit - //impl_i4be!(i16); - //impl_i4be!(u16); - // - //impl_i2be!(i8); - //impl_i4be!(i8); - //impl_i2be!(u8); - //impl_i4be!(u8); - } + + +#[no_mangle] +pub extern "C" fn f64_to_i16(arr: *mut f64, size: usize) -> f64 { impl_i16be!(f64, arr, size) } +#[no_mangle] +pub extern "C" fn f64_to_i32(arr: *mut f64, size: usize) -> f64 { impl_i32be!(f64, arr, size) } + +#[no_mangle] +pub extern "C" fn f32_to_i16(arr: *mut f32, size: usize) -> f32 { impl_i16be!(f32, arr, size) } +#[no_mangle] +pub extern "C" fn f32_to_i32(arr: *mut f32, size: usize) -> f32 { impl_i32be!(f32, arr, size) } + +#[no_mangle] +pub extern "C" fn u64_to_i16(arr: *mut u64, size: usize) -> u64 { impl_i16be!(u64, arr, size) } +#[no_mangle] +pub extern "C" fn u64_to_i32(arr: *mut u64, size: usize) -> u64 { impl_i32be!(u64, arr, size) } + +#[no_mangle] +pub extern "C" fn i64_to_i16(arr: *mut i64, size: usize) -> i64 { impl_i16be!(i64, arr, size) } +#[no_mangle] +pub extern "C" fn i64_to_i32(arr: *mut i64, size: usize) -> i64 { impl_i32be!(i64, arr, size) } + +#[no_mangle] +pub extern "C" fn u32_to_i16(arr: *mut u32, size: usize) -> u32 { impl_i16be!(u32, arr, size) } +#[no_mangle] +pub extern "C" fn u32_to_i32(arr: *mut u32, size: usize) -> u32 { impl_i32be!(u32, arr, size) } + +#[no_mangle] +pub extern "C" fn i32_to_i16(arr: *mut i32, size: usize) -> i32 { impl_i16be!(i32, arr, size) } +#[no_mangle] +pub extern "C" fn i32_to_i32(arr: *mut i32, size: usize) -> i32 { impl_i32be!(i32, arr, size) } + + +#[no_mangle] +pub extern "C" fn u16_to_i16(arr: *mut u16, size: usize) -> u16 { impl_i16be!(u16, arr, size) } + +#[no_mangle] +pub extern "C" fn i16_to_i16(arr: *mut i16, size: usize) -> i16 { impl_i16be!(i16, arr, size) } +