Introduction

Symjit is a lightweight just-in-time (JIT) optimizer compiler for mathematical expressions written in Rust. It was originally designed to compile SymPy (Python’s symbolic algebra package) expressions into machine code and to serve as a bridge between SymPy and numerical routines provided by NumPy and SciPy libraries.

Symjit emits AMD64 (x86-64), ARM64 (aarch64), and 64-bit RISC-V (riscv64) machine codes on Linux, Windows, and macOS platforms. SIMD is supported on x86-64 CPUs with AVX instruction sets.

Symbolica (https://symbolica.io/) is a fast Rust-based Computer Algebra System. Symbolica usually generate fast code using external compilers (e.g., using gcc to compile synthetic c++ code). Symjit accepts Symbolica expressions and can act as an optional code-generator for Symbolica.

Symjit-bridge crate acts as a bridge between Symbolica and Symjit to ease generating JIT code for Symbolica expressions.

Workflow

The main workflow is using different Runners. A runner corresponds to a Symbolica CompiledEvaluator object. The main runners are:

• CompiledRealRunner, corresponding to CompiledRealEvaluator.
• CompiledComplexRunner, corresponding to CompiledComplexEvaluator.
• InterpretedRealRunner, bytecode interpreter, generally similar to ExpressionEvaluator.
• InterpretedComplexRunner, bytecode interpreter, generally similar to ExpressionEvaluator.

Each runner has four main methods:

• compile(ev: &ExpressionEvaluator<T>, config: Config): the main constructor. T is either f64 or Complex<f64>, and config is an object of type Config. For most applications, the default config suffices. However, Config.use_threads(bool) is useful to enable multi-threading.
• compile_with_funcs(ev: &ExpressionEvaluator<T>, config: Config, df: &Defuns): Same as compile but with the additional of external functions defined in a Defuns structure.
• evaluate(args, outs): similar to the corresponding method of the Evaluators.
• save(filename).
• load(filename).

Both CompiledRealRunner and CompiledComplexRunner may use SIMD instructions if it is available and the number of input rows is equal or more than the number of SIMD lanes (4 in AVX, 2 in aarch64).

use anyhow::Result;
use symjit_bridge::{compile, Config};

use symbolica::{
    atom::AtomCore,
    evaluate::{FunctionMap, OptimizationSettings},
    parse, symbol,
};

fn test_real_runner() -> Result<()> {
    let params = vec![parse!("x"), parse!("y")];
    let f = FunctionMap::new();
    let ev = parse!("x + y^3")
        .evaluator(&f, &params, OptimizationSettings::default())
        .unwrap()
        .map_coeff(&|x| x.re.to_f64());

    let mut runner = CompiledRealRunner::compile(&ev, Config::default())?;
    let mut outs: [f64; 1] = [0.0];
    runner.evaluate(&[3.0, 5.0], &mut outs);
    assert_eq!(outs[0], 128.0);
    Ok(())
}

External Functions

Symjit has a rich set of transcendental, conditional, and logical functions (refer to Symjit for details). It is possible to expose these functions to Symbolica by using add_external_function:

fn test_external() -> Result<()> {
    let params = vec![parse!("x"), parse!("y")];
    let mut f = FunctionMap::new();
    f.add_external_function(symbol!("sinh"), "sinh".to_string())
        .unwrap();

    let ev = parse!("sinh(x+y)")
        .evaluator(&f, &params, OptimizationSettings::default())
        .unwrap()
        .map_coeff(&|x| Complex::new(x.re.to_f64(), x.im.to_f64()));

    let mut runner = CompiledComplexRunner::compile(&ev, Config::default())?;
    let args = [Complex::new(1.0, 2.0), Complex::new(2.0, -1.0)];
    let mut outs = [Complex::<f64>::default(); 1];
    runner.evaluate(&args, &mut outs);
    assert_eq!(outs[0], Complex::new(3.0, 1.0).sinh());
    Ok(())
}