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//! * `CompiledSimdRealRunner`, corresponding to `CompiledSimdRealEvaluator`.
29//! * `CompiledSimdComplexRunner`, corresponding to `CompiledSimdComplexEvaluator`.
30//! * `CompiledScatteredSimdRealRunner`, similar to `CompiledSimdRealRunner` but the
31//! data layout is similar to `CompiledRealRunner`.
32//! * `CompiledScatteredSimdComplexRunner`, similar to `CompiledSimdComplexRunner` but
33//! the data layout is similar to `CompiledComplexRunner`.
34//! * `InterpretedRealRunner`, bytecode interpreter, generally similar to `ExpressionEvaluator`.
35//! * `InterpretedComplexRunner`, bytecode interpreter, generally similar to `ExpressionEvaluator`.
36//!
37//! Each runner has four main methods:
38//!
39//! * `compile(ev: &ExpressionEvaluator<T>, config: Config)`: the main constructor. `T` is either `f64`
40//! or `Complex<f64>`, and `config` is an object of type `Config`. For most applications, the
41//! default config suffices. However, `Config.use_threads(bool)` is useful to enable multi-threading.
42//! * `evaluate(args, outs)`: similar to the corresponding method of the `Evaluator`s.
43//! * `save(filename)`.
44//! * `load(filename)`.
45//!
46
47//! ```rust
48//! use anyhow::Result;
49//! use symjit_bridge::{compile, Config};
50
51//! use symbolica::{
52//! atom::AtomCore,
53//! evaluate::{FunctionMap, OptimizationSettings},
54//! parse, symbol,
55//! };
56//!
57//! fn test_real_runner() -> Result<()> {
58//! let params = vec![parse!("x"), parse!("y")];
59//! let f = FunctionMap::new();
60//! let ev = parse!("x + y^3")
61//! .evaluator(&f, ¶ms, OptimizationSettings::default())
62//! .unwrap()
63//! .map_coeff(&|x| x.re.to_f64());
64
65//! let mut runner = CompiledRealRunner::compile(&ev, Config::default())?;
66//! let mut outs: [f64; 1] = [0.0];
67//! runner.evaluate(&[3.0, 5.0], &mut outs);
68//! assert_eq!(outs[0], 128.0);
69//! Ok(())
70//! }
71//! ```
72
73//! ## External Functions
74
75//! Symjit has a rich set of transcendental, conditional, and logical functions (refer to
76//! [Symjit](https://github.com/siravan/symjit) for details). It is possible to expose
77//! these functions to Symbolica by using `add_external_function`:
78
79//! ```rust
80//! fn test_external() -> Result<()> {
81//! let params = vec![parse!("x"), parse!("y")];
82//! let mut f = FunctionMap::new();
83//! f.add_external_function(symbol!("sinh"), "sinh".to_string())
84//! .unwrap();
85
86//! let ev = parse!("sinh(x+y)")
87//! .evaluator(&f, ¶ms, OptimizationSettings::default())
88//! .unwrap()
89//! .map_coeff(&|x| Complex::new(x.re.to_f64(), x.im.to_f64()));
90
91//! let mut runner = CompiledComplexRunner::compile(&ev, Config::default())?;
92//! let args = [Complex::new(1.0, 2.0), Complex::new(2.0, -1.0)];
93//! let mut outs = [Complex::<f64>::default(); 1];
94//! runner.evaluate(&args, &mut outs);
95//! assert_eq!(outs[0], Complex::new(3.0, 1.0).sinh());
96//! Ok(())
97//! }
98//! ```
99//!
100
101use anyhow::Result;
102use num_complex::Complex;
103
104pub use runners::{
105 CompiledComplexRunner, CompiledRealRunner, CompiledScatteredSimdComplexRunner,
106 CompiledScatteredSimdRealRunner, CompiledSimdComplexRunner, CompiledSimdRealRunner,
107 InterpretedComplexRunner, InterpretedRealRunner,
108};
109use symjit::{compiler, Translator};
110pub use symjit::{Application, Config};
111
112use symbolica::evaluate::{BuiltinSymbol, ExpressionEvaluator, Instruction, Slot};
113
114mod runners;
115
116fn slot(s: Slot) -> compiler::Slot {
117 match s {
118 Slot::Param(id) => compiler::Slot::Param(id),
119 Slot::Out(id) => compiler::Slot::Out(id),
120 Slot::Const(id) => compiler::Slot::Const(id),
121 Slot::Temp(id) => compiler::Slot::Temp(id),
122 }
123}
124
125fn slot_list(v: &[Slot]) -> Vec<compiler::Slot> {
126 v.iter().map(|s| slot(*s)).collect::<Vec<compiler::Slot>>()
127}
128
129fn builtin_symbol(s: BuiltinSymbol) -> compiler::BuiltinSymbol {
130 compiler::BuiltinSymbol(s.get_symbol().get_id())
131}
132
133fn translate(
134 instructions: Vec<Instruction>,
135 constants: Vec<Complex<f64>>,
136 config: Config,
137) -> Result<Translator> {
138 let mut translator = Translator::new(config);
139
140 for z in constants {
141 translator.append_constant(z)?;
142 }
143
144 for q in instructions {
145 match q {
146 Instruction::Add(lhs, args, num_reals) => {
147 translator.append_add(&slot(lhs), &slot_list(&args), num_reals)?
148 }
149 Instruction::Mul(lhs, args, num_reals) => {
150 translator.append_mul(&slot(lhs), &slot_list(&args), num_reals)?
151 }
152 Instruction::Pow(lhs, arg, p, is_real) => {
153 translator.append_pow(&slot(lhs), &slot(arg), p, is_real)?
154 }
155 Instruction::Powf(lhs, arg, p, is_real) => {
156 translator.append_powf(&slot(lhs), &slot(arg), &slot(p), is_real)?
157 }
158 Instruction::Assign(lhs, rhs) => translator.append_assign(&slot(lhs), &slot(rhs))?,
159 Instruction::Fun(lhs, fun, arg, is_real) => {
160 translator.append_fun(&slot(lhs), &builtin_symbol(fun), &slot(arg), is_real)?
161 }
162 Instruction::Join(lhs, cond, true_val, false_val) => translator.append_join(
163 &slot(lhs),
164 &slot(cond),
165 &slot(true_val),
166 &slot(false_val),
167 )?,
168 Instruction::Label(id) => translator.append_label(id)?,
169 Instruction::IfElse(cond, id) => translator.append_if_else(&slot(cond), id)?,
170 Instruction::Goto(id) => translator.append_goto(id)?,
171 Instruction::ExternalFun(lhs, op, args) => {
172 translator.append_external_fun(&slot(lhs), &op, &slot_list(&args))?
173 }
174 }
175 }
176
177 Ok(translator)
178}
179
180pub trait Number {
181 fn as_complex(&self) -> Complex<f64>;
182}
183
184impl Number for Complex<f64> {
185 fn as_complex(&self) -> Complex<f64> {
186 *self
187 }
188}
189
190impl Number for f64 {
191 fn as_complex(&self) -> Complex<f64> {
192 Complex::new(*self, 0.0)
193 }
194}
195
196pub fn compile<T: Clone + Number>(
197 ev: &ExpressionEvaluator<T>,
198 config: Config,
199) -> Result<Application> {
200 let (instructions, _, constants) = ev.export_instructions();
201 let constants: Vec<Complex<f64>> = constants.iter().map(|x| x.as_complex()).collect();
202 let mut translator = translate(instructions, constants, config).unwrap();
203 translator.compile()
204}