CFFI + maturin approach
This commit is contained in:
parent
320958d888
commit
bac31e2ce6
@ -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()
|
||||
|
||||
@ -31,3 +31,7 @@ dynamic = ["version"]
|
||||
dependencies = [
|
||||
"cffi",
|
||||
]
|
||||
|
||||
[tool.maturin]
|
||||
bindings = "cffi"
|
||||
|
||||
|
||||
157
src/lib.rs
157
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<Self, Self> where Self: Sized;
|
||||
}
|
||||
|
||||
pub trait ToInt4BE {
|
||||
fn convert_to_i4be(ii: Self) -> Result<Self, Self> 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) }
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user