Skip to main content

ronn/
lib.rs

1#![allow(dead_code)]
2#![allow(non_local_definitions)]
3//! Python bindings for RONN
4//!
5//! Provides a Pythonic interface to RONN's high-performance ONNX inference runtime.
6//!
7//! # Installation
8//!
9//! ```bash
10//! pip install ronn
11//! ```
12//!
13//! # Quick Start
14//!
15//! ```python
16//! import ronn
17//! import numpy as np
18//!
19//! # Load model
20//! model = ronn.Model.load("model.onnx")
21//!
22//! # Create session
23//! session = model.create_session(
24//!     optimization_level="O3",
25//!     provider="cpu"
26//! )
27//!
28//! # Run inference
29//! inputs = {"input": np.array([[1.0, 2.0, 3.0]], dtype=np.float32)}
30//! outputs = session.run(inputs)
31//! print(outputs["output"])
32//! ```
33
34use pyo3::prelude::*;
35
36mod error;
37mod model;
38mod session;
39mod tensor;
40
41pub use error::RonnError;
42pub use model::PyModel;
43pub use session::PySession;
44pub use tensor::PyTensor;
45
46/// RONN Python module
47#[pymodule]
48fn ronn(_py: Python, m: &PyModule) -> PyResult<()> {
49    m.add_class::<PyModel>()?;
50    m.add_class::<PySession>()?;
51    m.add_class::<PyTensor>()?;
52    m.add_class::<PyOptimizationLevel>()?;
53    m.add_class::<PyProviderType>()?;
54    m.add_class::<PyBatchConfig>()?;
55
56    // Add version
57    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
58
59    Ok(())
60}
61
62/// Optimization level for graph optimizations
63#[pyclass(name = "OptimizationLevel")]
64#[derive(Clone)]
65pub struct PyOptimizationLevel {
66    inner: ronn_core::OptimizationLevel,
67}
68
69#[pymethods]
70impl PyOptimizationLevel {
71    /// No optimizations (O0)
72    #[staticmethod]
73    fn none() -> Self {
74        Self {
75            inner: ronn_core::OptimizationLevel::None,
76        }
77    }
78
79    /// Basic optimizations (O1)
80    #[staticmethod]
81    fn basic() -> Self {
82        Self {
83            inner: ronn_core::OptimizationLevel::Basic,
84        }
85    }
86
87    /// Default optimizations (O2, recommended)
88    #[staticmethod]
89    fn default() -> Self {
90        Self {
91            inner: ronn_core::OptimizationLevel::Basic,
92        }
93    }
94
95    /// Aggressive optimizations (O3)
96    #[staticmethod]
97    fn aggressive() -> Self {
98        Self {
99            inner: ronn_core::OptimizationLevel::Aggressive,
100        }
101    }
102}
103
104impl Default for PyOptimizationLevel {
105    fn default() -> Self {
106        Self::default()
107    }
108}
109
110/// Execution provider type
111#[pyclass(name = "ProviderType")]
112#[derive(Clone)]
113pub struct PyProviderType {
114    inner: ronn_core::ProviderId,
115}
116
117#[pymethods]
118impl PyProviderType {
119    /// CPU provider (default)
120    #[staticmethod]
121    fn cpu() -> Self {
122        Self {
123            inner: ronn_core::ProviderId::CPU,
124        }
125    }
126
127    /// GPU provider (CUDA)
128    #[staticmethod]
129    fn gpu() -> Self {
130        Self {
131            inner: ronn_core::ProviderId::GPU,
132        }
133    }
134
135    /// BitNet provider (1.58-bit quantization)
136    #[staticmethod]
137    fn bitnet() -> Self {
138        Self {
139            inner: ronn_core::ProviderId::BitNet,
140        }
141    }
142
143    /// WebAssembly provider
144    #[staticmethod]
145    fn wasm() -> Self {
146        Self {
147            inner: ronn_core::ProviderId::WebAssembly,
148        }
149    }
150}
151
152impl Default for PyProviderType {
153    fn default() -> Self {
154        Self::cpu()
155    }
156}
157
158/// Batch processing configuration
159#[pyclass(name = "BatchConfig")]
160#[derive(Clone)]
161pub struct PyBatchConfig {
162    #[pyo3(get, set)]
163    /// Maximum batch size
164    pub max_batch_size: usize,
165
166    #[pyo3(get, set)]
167    /// Timeout in milliseconds
168    pub timeout_ms: u64,
169
170    #[pyo3(get, set)]
171    /// Queue capacity
172    pub queue_capacity: usize,
173}
174
175#[pymethods]
176impl PyBatchConfig {
177    #[new]
178    #[pyo3(signature = (max_batch_size=32, timeout_ms=10, queue_capacity=1024))]
179    fn new(max_batch_size: usize, timeout_ms: u64, queue_capacity: usize) -> Self {
180        Self {
181            max_batch_size,
182            timeout_ms,
183            queue_capacity,
184        }
185    }
186}
187
188impl Default for PyBatchConfig {
189    fn default() -> Self {
190        Self::new(32, 10, 1024)
191    }
192}