symjit_bridge/
lib.rs

1//! # Introduction
2
3//! [Symjit](https://github.com/siravan/symjit) is a lightweight just-in-time (JIT)
4//! optimizer compiler for mathematical expressions written in Rust. It was originally
5//! designed to compile SymPy (Python’s symbolic algebra package) expressions
6//! into machine code and to serve as a bridge between SymPy and numerical routines
7//! provided by NumPy and SciPy libraries.
8
9//! Symjit emits AMD64 (x86-64), ARM64 (aarch64), and 64-bit RISC-V (riscv64) machine
10//! codes on Linux, Windows, and macOS platforms. SIMD is supported on x86-64
11//! CPUs with AVX instruction sets.
12
13//! Symbolica (<https://symbolica.io/>) is a fast Rust-based Computer Algebra System.
14//! Symbolica usually generate fast code using external compilers (e.g., using gcc to
15//! compile synthetic c++ code). Symjit accepts Symbolica expressions and can act as
16//! an optional code-generator for Symbolica.
17
18//! Symjit-bridge crate acts as a bridge between Symbolica and Symjit to ease generating
19//! JIT code for Symbolica expressions.
20
21//! # Workflow
22
23//! The main workflow is using different `Runner`s. A runner corresponds to a Symbolica
24//! `CompiledEvaluator` object. The main runners are:
25//!
26//! * `CompiledRealRunner`, corresponding to `CompiledRealEvaluator`.
27//! * `CompiledComplexRunner`, corresponding to `CompiledComplexEvaluator`.
28//! * `InterpretedRealRunner`, bytecode interpreter, generally similar to `ExpressionEvaluator`.
29//! * `InterpretedComplexRunner`, bytecode interpreter, generally similar to `ExpressionEvaluator`.
30//!
31//! Each runner has four main methods:
32//!
33//! * `compile(ev: &ExpressionEvaluator<T>, config: Config)`: the main constructor. `T` is either `f64`
34//!     or `Complex<f64>`, and `config` is an object of type `Config`. For most applications, the
35//!     default config suffices. However, `Config.use_threads(bool)` is useful to enable multi-threading.
36//! * `compile_with_funcs(ev: &ExpressionEvaluator<T>, config: Config, df: &Defuns, num_params: usize)`: Same as
37//!     `compile` but with the additional of external functions defined in a `Defuns` structure and `num_prams`.
38//! * `compile_string(model: String, config: Config)`: `model` is a string generated using `get_instruction` method
39//!     in Python, and `config` is an object of type `Config`.
40//! * `compile_string_with_funcs(model: String, config: Config, df: &Defuns, num_params: usize)`: Same as
41//!     `compile_string` but with the additional of external functions defined in a `Defuns` structure
42//!     and `num_params`.
43//! * `evaluate(args, outs)`: similar to the corresponding method of the `Evaluator`s.
44//! * `save(filename)`.
45//! * `load(filename)`.
46//!
47//! Both `CompiledRealRunner` and `CompiledComplexRunner` may use SIMD instructions if it is available
48//!     and the number of input rows is equal or more than the number of SIMD lanes (4 in AVX, 2 in aarch64).
49//!
50//! ```rust
51//! use anyhow::Result;
52//! use symjit_bridge::{compile, Config};
53
54//! use symbolica::{
55//!     atom::AtomCore,
56//!     evaluate::{FunctionMap, OptimizationSettings},
57//!     parse, symbol,
58//! };
59//!
60//! fn test_real_runner() -> Result<()> {
61//!     let params = vec![parse!("x"), parse!("y")];
62//!     let f = FunctionMap::new();
63//!     let ev = parse!("x + y^3")
64//!         .evaluator(&f, &params, OptimizationSettings::default())
65//!         .unwrap()
66//!         .map_coeff(&|x| x.re.to_f64());
67
68//!     let mut runner = CompiledRealRunner::compile(&ev, Config::default())?;
69//!     let mut outs: [f64; 1] = [0.0];
70//!     runner.evaluate(&[3.0, 5.0], &mut outs);
71//!     assert_eq!(outs[0], 128.0);
72//!     Ok(())
73//! }
74//! ```
75
76//! ## External Functions
77
78//! Symjit has a rich set of transcendental, conditional, and logical functions (refer to
79//! [Symjit](https://github.com/siravan/symjit) for details). It is possible to expose
80//! these functions to Symbolica by using `add_external_function`:
81
82//! ```rust
83//! fn test_external() -> Result<()> {
84//!     let params = vec![parse!("x"), parse!("y")];
85//!     let mut f = FunctionMap::new();
86//!     f.add_external_function(symbol!("sinh"), "sinh".to_string())
87//!         .unwrap();
88
89//!     let ev = parse!("sinh(x+y)")
90//!         .evaluator(&f, &params, OptimizationSettings::default())
91//!         .unwrap()
92//!         .map_coeff(&|x| Complex::new(x.re.to_f64(), x.im.to_f64()));
93
94//!     let mut runner = CompiledComplexRunner::compile(&ev, Config::default())?;
95//!     let args = [Complex::new(1.0, 2.0), Complex::new(2.0, -1.0)];
96//!     let mut outs = [Complex::<f64>::default(); 1];
97//!     runner.evaluate(&args, &mut outs);
98//!     assert_eq!(outs[0], Complex::new(3.0, 1.0).sinh());
99//!     Ok(())
100//! }
101//! ```
102//!
103
104use anyhow::Result;
105
106pub use runners::{
107    CompiledComplexRunner, CompiledRealRunner, InterpretedComplexRunner, InterpretedRealRunner,
108};
109use symjit::{instruction, Compiler, Composer, Translator, Transliterator};
110pub use symjit::{Application, Complex, ComplexFloat, Config, Defuns};
111
112use symbolica::evaluate::{BuiltinSymbol, ExpressionEvaluator, Instruction, Slot};
113
114mod runners;
115
116fn slot(s: Slot) -> instruction::Slot {
117    match s {
118        Slot::Param(id) => instruction::Slot::Param(id),
119        Slot::Out(id) => instruction::Slot::Out(id),
120        Slot::Const(id) => instruction::Slot::Const(id),
121        Slot::Temp(id) => instruction::Slot::Temp(id),
122    }
123}
124
125fn slot_list(v: &[Slot]) -> Vec<instruction::Slot> {
126    v.iter()
127        .map(|s| slot(*s))
128        .collect::<Vec<instruction::Slot>>()
129}
130
131fn builtin_symbol(s: BuiltinSymbol) -> instruction::BuiltinSymbol {
132    instruction::BuiltinSymbol(s.get_symbol().get_id())
133}
134
135fn translate(
136    instructions: Vec<Instruction>,
137    constants: Vec<Complex<f64>>,
138    config: Config,
139    df: Defuns,
140    direct: bool,
141) -> Result<Box<dyn Composer>> {
142    let mut translator: Box<dyn Composer> = if direct {
143        Box::new(Transliterator::new(config, df))
144    } else {
145        Box::new(Translator::new(config, df))
146    };
147
148    for z in constants {
149        translator.append_constant(z)?;
150    }
151
152    for q in instructions {
153        match q {
154            Instruction::Add(lhs, args, num_reals) => {
155                translator.append_add(&slot(lhs), &slot_list(&args), num_reals)?
156            }
157            Instruction::Mul(lhs, args, num_reals) => {
158                translator.append_mul(&slot(lhs), &slot_list(&args), num_reals)?
159            }
160            Instruction::Pow(lhs, arg, p, is_real) => {
161                translator.append_pow(&slot(lhs), &slot(arg), p, is_real)?
162            }
163            Instruction::Powf(lhs, arg, p, is_real) => {
164                translator.append_powf(&slot(lhs), &slot(arg), &slot(p), is_real)?
165            }
166            Instruction::Assign(lhs, rhs) => translator.append_assign(&slot(lhs), &slot(rhs))?,
167            Instruction::Fun(lhs, fun, arg, is_real) => {
168                translator.append_fun(&slot(lhs), &builtin_symbol(fun), &slot(arg), is_real)?
169            }
170            Instruction::Join(lhs, cond, true_val, false_val) => translator.append_join(
171                &slot(lhs),
172                &slot(cond),
173                &slot(true_val),
174                &slot(false_val),
175            )?,
176            Instruction::Label(id) => translator.append_label(id)?,
177            Instruction::IfElse(cond, id) => translator.append_if_else(&slot(cond), id)?,
178            Instruction::Goto(id) => translator.append_goto(id)?,
179            Instruction::ExternalFun(lhs, op, args) => {
180                translator.append_external_fun(&slot(lhs), &op, &slot_list(&args))?
181            }
182        }
183    }
184
185    Ok(translator)
186}
187
188pub trait Number {
189    fn as_complex(&self) -> Complex<f64>;
190}
191
192impl Number for Complex<f64> {
193    fn as_complex(&self) -> Complex<f64> {
194        *self
195    }
196}
197
198impl Number for f64 {
199    fn as_complex(&self) -> Complex<f64> {
200        Complex::new(*self, 0.0)
201    }
202}
203
204pub fn compile<T: Clone + Number>(
205    ev: &ExpressionEvaluator<T>,
206    config: Config,
207    df: Defuns,
208    num_params: usize,
209) -> Result<Application> {
210    let (instructions, _, constants) = ev.export_instructions();
211    let constants: Vec<Complex<f64>> = constants.iter().map(|x| x.as_complex()).collect();
212    let mut translator = translate(instructions, constants, config, df, true).unwrap();
213    translator.set_num_params(num_params);
214    translator.compile()
215}
216
217pub fn compile_string(
218    model: String,
219    config: Config,
220    df: Defuns,
221    num_params: usize,
222) -> Result<Application> {
223    let mut comp = Compiler::with_config(config);
224    comp.translate(model, df, num_params)
225}
symjit_bridge/lib.rs

symjit_bridge/
lib.rs
