from __future__ import annotations
import ctypes
import os
import platform
from dataclasses import dataclass
from pathlib import Path
try:
from .errors import InvalidInputError, RnnFfiError
except ImportError:
from errors import InvalidInputError, RnnFfiError
class Codes:
OK = 0
class _RnnFfiBenchmarkRecord(ctypes.Structure):
_fields_ = [
("model_name_ptr", ctypes.POINTER(ctypes.c_uint8)),
("model_name_len", ctypes.c_size_t),
("precision_ptr", ctypes.POINTER(ctypes.c_uint8)),
("precision_len", ctypes.c_size_t),
("elapsed_ms", ctypes.c_uint64),
("iterations", ctypes.c_uint64),
("avg_loss", ctypes.c_float),
("last_loss", ctypes.c_float),
("output_bytes", ctypes.c_uint64),
("train_samples", ctypes.c_uint64),
("total_params", ctypes.c_uint64),
("layer_count", ctypes.c_uint32),
("input_dim", ctypes.c_uint32),
("output_dim", ctypes.c_uint32),
("benchmark_flags", ctypes.c_uint64),
("weights_bytes", ctypes.c_uint64),
("biases_bytes", ctypes.c_uint64),
("min_loss", ctypes.c_float),
("max_loss", ctypes.c_float),
("loss_stddev", ctypes.c_float),
("iterations_per_sec", ctypes.c_float),
("samples_per_sec", ctypes.c_float),
]
class _RnnFfiBenchmarkView(ctypes.Structure):
_fields_ = [
("model_name_ptr", ctypes.POINTER(ctypes.c_uint8)),
("model_name_len", ctypes.c_size_t),
("precision_ptr", ctypes.POINTER(ctypes.c_uint8)),
("precision_len", ctypes.c_size_t),
("elapsed_ms", ctypes.c_uint64),
("iterations", ctypes.c_uint64),
("avg_loss", ctypes.c_float),
("last_loss", ctypes.c_float),
("output_bytes", ctypes.c_uint64),
("train_samples", ctypes.c_uint64),
("total_params", ctypes.c_uint64),
("layer_count", ctypes.c_uint32),
("input_dim", ctypes.c_uint32),
("output_dim", ctypes.c_uint32),
("benchmark_flags", ctypes.c_uint64),
("weights_bytes", ctypes.c_uint64),
("biases_bytes", ctypes.c_uint64),
("min_loss", ctypes.c_float),
("max_loss", ctypes.c_float),
("loss_stddev", ctypes.c_float),
("iterations_per_sec", ctypes.c_float),
("samples_per_sec", ctypes.c_float),
]
@dataclass(frozen=True)
class BenchmarkRecord:
model_name: str
precision: str
elapsed_ms: int
iterations: int
avg_loss: float
last_loss: float
output_bytes: int
train_samples: int = 0
total_params: int = 0
layer_count: int = 0
input_dim: int = 0
output_dim: int = 0
benchmark_flags: int = 0
weights_bytes: int = 0
biases_bytes: int = 0
min_loss: float = 0.0
max_loss: float = 0.0
loss_stddev: float = 0.0
iterations_per_sec: float = 0.0
samples_per_sec: float = 0.0
@dataclass(frozen=True)
class BenchmarkView:
model_name: str
precision: str
elapsed_ms: int
iterations: int
avg_loss: float
last_loss: float
output_bytes: int
train_samples: int
total_params: int
layer_count: int
input_dim: int
output_dim: int
benchmark_flags: int
weights_bytes: int
biases_bytes: int
min_loss: float
max_loss: float
loss_stddev: float
iterations_per_sec: float
samples_per_sec: float
def _candidate_library_names() -> list[str]:
system = platform.system().lower()
if "windows" in system:
return ["rnn.dll", "native_neural_network.dll"]
if "darwin" in system or "mac" in system:
return ["librnn.dylib", "rnn.dylib", "libnative_neural_network.dylib"]
return ["librnn.so", "rnn.so", "libnative_neural_network.so"]
def _discover_library_path(explicit: str | None = None) -> str:
if explicit:
return explicit
env = os.getenv("RNN_FFI_LIB")
if env:
return env
repo_root = Path(__file__).resolve().parents[2]
target_candidates = [repo_root / "target" / "debug", repo_root / "target" / "release"]
searched: list[str] = []
for folder in target_candidates:
for name in _candidate_library_names():
path = folder / name
searched.append(str(path))
if path.exists():
return str(path)
names = ", ".join(_candidate_library_names())
hint = (
"Native FFI library not found. "
f"Expected one of [{names}] under target/debug or target/release, "
"or set RNN_FFI_LIB to an explicit library path."
)
details = "\n".join(searched)
raise FileNotFoundError(f"{hint}\nSearched:\n{details}")
class FfiLibrary:
def __init__(self, library_path: str | None = None):
lib_path = _discover_library_path(library_path)
self._lib = ctypes.CDLL(lib_path)
self._configure_symbols()
def _configure_symbols(self) -> None:
self._lib.rnn_ffi_api_version.argtypes = []
self._lib.rnn_ffi_api_version.restype = ctypes.c_uint32
self._lib.rnn_ffi_benchmark_encoded_size.argtypes = [
ctypes.POINTER(_RnnFfiBenchmarkRecord),
ctypes.POINTER(ctypes.c_size_t),
]
self._lib.rnn_ffi_benchmark_encoded_size.restype = ctypes.c_int32
self._lib.rnn_ffi_encode_benchmark_blob.argtypes = [
ctypes.POINTER(_RnnFfiBenchmarkRecord),
ctypes.POINTER(ctypes.c_uint8),
ctypes.c_size_t,
ctypes.POINTER(ctypes.c_size_t),
]
self._lib.rnn_ffi_encode_benchmark_blob.restype = ctypes.c_int32
self._lib.rnn_ffi_decode_benchmark_blob.argtypes = [
ctypes.POINTER(ctypes.c_uint8),
ctypes.c_size_t,
ctypes.POINTER(_RnnFfiBenchmarkView),
]
self._lib.rnn_ffi_decode_benchmark_blob.restype = ctypes.c_int32
self._lib.rnn_ffi_error_message.argtypes = [ctypes.c_int32]
self._lib.rnn_ffi_error_message.restype = ctypes.c_char_p
def api_version(self) -> int:
return int(self._lib.rnn_ffi_api_version())
def error_message(self, code: int) -> str:
raw = self._lib.rnn_ffi_error_message(ctypes.c_int32(code))
if not raw:
return "unknown error"
return raw.decode("utf-8", errors="replace")
def _raise_if_failed(self, code: int) -> None:
if code != Codes.OK:
raise RnnFfiError(code, self.error_message(code))
def _build_record(self, record: BenchmarkRecord):
if not record.model_name or not record.precision:
raise InvalidInputError("model_name and precision must be non-empty")
model_bytes = record.model_name.encode("utf-8")
precision_bytes = record.precision.encode("utf-8")
model_buf = (ctypes.c_uint8 * len(model_bytes)).from_buffer_copy(model_bytes)
precision_buf = (ctypes.c_uint8 * len(precision_bytes)).from_buffer_copy(precision_bytes)
ffi_record = _RnnFfiBenchmarkRecord(
model_name_ptr=ctypes.cast(model_buf, ctypes.POINTER(ctypes.c_uint8)),
model_name_len=len(model_bytes),
precision_ptr=ctypes.cast(precision_buf, ctypes.POINTER(ctypes.c_uint8)),
precision_len=len(precision_bytes),
elapsed_ms=int(record.elapsed_ms),
iterations=int(record.iterations),
avg_loss=float(record.avg_loss),
last_loss=float(record.last_loss),
output_bytes=int(record.output_bytes),
train_samples=int(record.train_samples),
total_params=int(record.total_params),
layer_count=int(record.layer_count),
input_dim=int(record.input_dim),
output_dim=int(record.output_dim),
benchmark_flags=int(record.benchmark_flags),
weights_bytes=int(record.weights_bytes),
biases_bytes=int(record.biases_bytes),
min_loss=float(record.min_loss),
max_loss=float(record.max_loss),
loss_stddev=float(record.loss_stddev),
iterations_per_sec=float(record.iterations_per_sec),
samples_per_sec=float(record.samples_per_sec),
)
return ffi_record
def encode_benchmark_blob(self, record: BenchmarkRecord) -> bytes:
ffi_record = self._build_record(record)
out_size = ctypes.c_size_t(0)
code = int(self._lib.rnn_ffi_benchmark_encoded_size(ctypes.byref(ffi_record), ctypes.byref(out_size)))
self._raise_if_failed(code)
out_buf = (ctypes.c_uint8 * out_size.value)()
out_used = ctypes.c_size_t(0)
code = int(
self._lib.rnn_ffi_encode_benchmark_blob(
ctypes.byref(ffi_record),
out_buf,
out_size.value,
ctypes.byref(out_used),
)
)
self._raise_if_failed(code)
return bytes(bytearray(out_buf)[: out_used.value])
def decode_benchmark_blob(self, blob: bytes) -> BenchmarkView:
if not blob:
raise InvalidInputError("blob cannot be empty")
blob_arr = (ctypes.c_uint8 * len(blob)).from_buffer_copy(blob)
out = _RnnFfiBenchmarkView()
code = int(self._lib.rnn_ffi_decode_benchmark_blob(blob_arr, len(blob), ctypes.byref(out)))
self._raise_if_failed(code)
model_bytes = ctypes.string_at(out.model_name_ptr, out.model_name_len)
precision_bytes = ctypes.string_at(out.precision_ptr, out.precision_len)
return BenchmarkView(
model_name=model_bytes.decode("utf-8", errors="strict"),
precision=precision_bytes.decode("utf-8", errors="strict"),
elapsed_ms=int(out.elapsed_ms),
iterations=int(out.iterations),
avg_loss=float(out.avg_loss),
last_loss=float(out.last_loss),
output_bytes=int(out.output_bytes),
train_samples=int(out.train_samples),
total_params=int(out.total_params),
layer_count=int(out.layer_count),
input_dim=int(out.input_dim),
output_dim=int(out.output_dim),
benchmark_flags=int(out.benchmark_flags),
weights_bytes=int(out.weights_bytes),
biases_bytes=int(out.biases_bytes),
min_loss=float(out.min_loss),
max_loss=float(out.max_loss),
loss_stddev=float(out.loss_stddev),
iterations_per_sec=float(out.iterations_per_sec),
samples_per_sec=float(out.samples_per_sec),
)