add cranelift codegen backend

Cargo.toml

@@ -8,17 +8,17 @@ exclude = ["Enzyme/enzyme/benchmarks"]
 authors = ["Martin Robinson <[email protected]>"]
 repository = "https://github.com/martinjrobins/diffsl"
 
-[[bin]]
-name = "diffsl"
-
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-llvm13-0 = ["inkwell-130", "llvm-sys-130"]
-llvm14-0 = ["inkwell-140", "llvm-sys-140"]
-llvm15-0 = ["inkwell-150", "llvm-sys-150"]
-llvm16-0 = ["inkwell-160", "llvm-sys-160"]
-llvm17-0 = ["inkwell-170", "llvm-sys-170"]
+llvm13-0 = ["inkwell-130", "llvm-sys-130", "inkwell_internals", "llvm", "enzyme"]
+llvm14-0 = ["inkwell-140", "llvm-sys-140", "inkwell_internals", "llvm", "enzyme"]
+llvm15-0 = ["inkwell-150", "llvm-sys-150", "inkwell_internals", "llvm", "enzyme"]
+llvm16-0 = ["inkwell-160", "llvm-sys-160", "inkwell_internals", "llvm", "enzyme"]
+llvm17-0 = ["inkwell-170", "llvm-sys-170", "inkwell_internals", "llvm", "enzyme"]
+llvm18-0 = ["inkwell-180", "llvm-sys-180", "inkwell_internals", "llvm", "enzyme"]
+enzyme = ["bindgen", "cmake"]
+llvm = []
 
 [dependencies]
 ndarray = { version = ">=0.15.0", features = ["approx-0_5"] } 
@@ -27,24 +27,31 @@ approx = ">=0.5"
 pest = ">=2.1.3"
 pest_derive = ">=2.1.0"
 itertools = ">=0.10.3"
-ouroboros = ">=0.17"
 clap = { version = "4.3.23", features = ["derive"] }
 uid = "0.1.7"
-inkwell-130 = { package = "inkwell", version = ">=0.4.0", features = ["llvm13-0"], optional = true }
-inkwell-140 = { package = "inkwell", version = ">=0.4.0", features = ["llvm14-0"], optional = true }
-inkwell-150 = { package = "inkwell", version = ">=0.4.0", features = ["llvm15-0"], optional = true }
-inkwell-160 = { package = "inkwell", version = ">=0.4.0", features = ["llvm16-0"], optional = true }
-inkwell-170 = { package = "inkwell", version = ">=0.4.0", features = ["llvm17-0"], optional = true }
+inkwell-130 = { package = "inkwell", version = ">=0.5.0", features = ["llvm13-0"], optional = true }
+inkwell-140 = { package = "inkwell", version = ">=0.5.0", features = ["llvm14-0"], optional = true }
+inkwell-150 = { package = "inkwell", version = ">=0.5.0", features = ["llvm15-0"], optional = true }
+inkwell-160 = { package = "inkwell", version = ">=0.5.0", features = ["llvm16-0"], optional = true }
+inkwell-170 = { package = "inkwell", version = ">=0.5.0", features = ["llvm17-0"], optional = true }
+inkwell-180 = { package = "inkwell", version = ">=0.5.0", features = ["llvm18-0"], optional = true }
 llvm-sys-130 = { package = "llvm-sys", version = "130.0.4", optional = true }
 llvm-sys-140 = { package = "llvm-sys", version = "140.0.2", optional = true }
 llvm-sys-150 = { package = "llvm-sys", version = "150.0.3", optional = true }
 llvm-sys-160 = { package = "llvm-sys", version = "160.1.0", optional = true }
 llvm-sys-170 = { package = "llvm-sys", version = "170.0.1", optional = true }
-inkwell_internals = "0.9.0"
+llvm-sys-180 = { package = "llvm-sys", version = "180.0.0", optional = true }
+inkwell_internals = { version = "0.9.0", optional = true }
+cranelift = "0.110.1"
+cranelift-module = "0.110.1"
+cranelift-jit = "0.110.1"
+cranelift-native = "0.110.1"
+target-lexicon = "0.12.16"
+aliasable = "0.1.3"
 
 [build-dependencies]
-cmake = "0.1.50"
-bindgen = "0.69.4"
+bindgen = { version = "0.69.4", optional = true }
+cmake = { version = "0.1.50", optional = true }
 
 [dev-dependencies]
 divan = "0.1.14"

benches/evaluation.rs

@@ -1,12 +1,15 @@
-use diffsl::{discretise::DiscreteModel, execution::Compiler, parser::parse_ds_string};
+use diffsl::{
+    discretise::DiscreteModel, execution::module::CodegenModule, parser::parse_ds_string, Compiler,
+    CraneliftModule,
+};
 use divan::Bencher;
 use ndarray::Array1;
 
 fn main() {
     divan::main();
 }
 
-fn setup(n: usize, f_text: &str, name: &str) -> Compiler {
+fn setup<M: CodegenModule>(n: usize, f_text: &str, name: &str) -> Compiler<M> {
     let u = vec![1.0; n];
     let full_text = format!(
         "
@@ -28,14 +31,32 @@ fn setup(n: usize, f_text: &str, name: &str) -> Compiler {
     );
     let model = parse_ds_string(&full_text).unwrap();
     let discrete_model = DiscreteModel::build(name, &model).unwrap();
-    let out = format!("test_output/benches_evaluation_{}", name);
-    Compiler::from_discrete_model(&discrete_model, out.as_str()).unwrap()
+    Compiler::from_discrete_model(&discrete_model).unwrap()
 }
 
+#[cfg(feature = "llvm")]
 #[divan::bench(consts = [1, 10, 100, 1000])]
-fn add_scalar_diffsl<const N: usize>(bencher: Bencher) {
+fn add_scalar_diffsl_llvm<const N: usize>(bencher: Bencher) {
+    use diffsl::LlvmModule;
+
+    let n = N;
+    let compiler = setup::<LlvmModule>(n, "u_i + 1.0", "add_scalar");
+    let mut data = compiler.get_new_data();
+    compiler.set_inputs(&[], data.as_mut_slice());
+    let mut u = vec![1.0; n];
+    compiler.set_u0(u.as_mut_slice(), data.as_mut_slice());
+    let mut rr = vec![0.0; n];
+    let t = 0.0;
+
+    bencher.bench_local(|| {
+        compiler.rhs(t, &u, &mut data, &mut rr);
+    });
+}
+
+#[divan::bench(consts = [1, 10, 100, 1000])]
+fn add_scalar_diffsl_cranelift<const N: usize>(bencher: Bencher) {
     let n = N;
-    let compiler = setup(n, "u_i + 1.0", "add_scalar");
+    let compiler = setup::<CraneliftModule>(n, "u_i + 1.0", "add_scalar");
     let mut data = compiler.get_new_data();
     compiler.set_inputs(&[], data.as_mut_slice());
     let mut u = vec![1.0; n];
@@ -47,6 +68,7 @@ fn add_scalar_diffsl<const N: usize>(bencher: Bencher) {
         compiler.rhs(t, &u, &mut data, &mut rr);
     });
 }
+
 #[divan::bench(consts = [1, 10, 100, 1000])]
 fn add_scalar_ndarray<const N: usize>(bencher: Bencher) {
     let n = N;

build.rs

@@ -1,74 +1,83 @@
-use bindgen::{BindgenError, Bindings, Builder};
-use std::{env, path::PathBuf};
+#[cfg(feature = "enzyme")]
+mod enzyme {
+    use bindgen::{BindgenError, Bindings, Builder};
+    use std::{env, path::PathBuf};
 
-fn compile_enzyme(llvm_dir: String) -> (String, String) {
-    let dst = cmake::Config::new("Enzyme/enzyme")
-        .define("ENZYME_STATIC_LIB", "ON")
-        .define("ENZYME_CLANG", "OFF")
-        .define("LLVM_DIR", llvm_dir)
-        .define(
-            "CMAKE_CXX_FLAGS",
-            "-Wno-comment -Wno-deprecated-declarations",
-        )
-        .build();
-    let dst_disp = dst.display();
-    let lib_dir = format!("{}/lib", dst_disp);
-    let inc_dir = "Enzyme/enzyme".to_string();
-    (lib_dir, inc_dir)
-}
+    fn compile_enzyme(llvm_dir: String) -> (String, String) {
+        let dst = cmake::Config::new("Enzyme/enzyme")
+            .define("ENZYME_STATIC_LIB", "ON")
+            .define("ENZYME_CLANG", "OFF")
+            .define("LLVM_DIR", llvm_dir)
+            .define(
+                "CMAKE_CXX_FLAGS",
+                "-Wno-comment -Wno-deprecated-declarations",
+            )
+            .build();
+        let dst_disp = dst.display();
+        let lib_dir = format!("{}/lib", dst_disp);
+        let inc_dir = "Enzyme/enzyme".to_string();
+        (lib_dir, inc_dir)
+    }
 
-fn enzyme_bindings(inc_dirs: &[String]) -> Result<Bindings, BindgenError> {
-    let mut builder = Builder::default()
-        .header("wrapper.h")
-        .generate_comments(false)
-        .clang_arg("-x")
-        .clang_arg("c++");
+    fn enzyme_bindings(inc_dirs: &[String]) -> Result<Bindings, BindgenError> {
+        let mut builder = Builder::default()
+            .header("wrapper.h")
+            .generate_comments(false)
+            .clang_arg("-x")
+            .clang_arg("c++");
 
-    // add include dirs
-    for dir in inc_dirs {
-        builder = builder.clang_arg(format!("-I{}", dir))
+        // add include dirs
+        for dir in inc_dirs {
+            builder = builder.clang_arg(format!("-I{}", dir))
+        }
+        builder.generate()
     }
-    builder.generate()
-}
 
-fn main() {
-    // get env vars matching DEP_LLVM_*_LIBDIR regex
-    let llvm_dirs: Vec<_> = env::vars()
-        .filter(|(k, _)| k.starts_with("DEP_LLVM_") && k.ends_with("_LIBDIR"))
-        .collect();
-    // take first one
-    let llvm_lib_dir = llvm_dirs
-        .first()
-        .expect("DEP_LLVM_*_LIBDIR not set")
-        .1
-        .clone();
-    let llvm_env_key = llvm_dirs.first().unwrap().0.clone();
-    let llvm_version = &llvm_env_key["DEP_LLVM_".len()..(llvm_env_key.len() - "_LIBDIR".len())];
-    dbg!(llvm_version);
+    pub fn enzyme_main() {
+        // get env vars matching DEP_LLVM_*_LIBDIR regex
+        let llvm_dirs: Vec<_> = env::vars()
+            .filter(|(k, _)| k.starts_with("DEP_LLVM_") && k.ends_with("_LIBDIR"))
+            .collect();
+        // take first one
+        let llvm_lib_dir = llvm_dirs
+            .first()
+            .expect("DEP_LLVM_*_LIBDIR not set")
+            .1
+            .clone();
+        let llvm_env_key = llvm_dirs.first().unwrap().0.clone();
+        let llvm_version = &llvm_env_key["DEP_LLVM_".len()..(llvm_env_key.len() - "_LIBDIR".len())];
+        dbg!(llvm_version);
 
-    // replace last "lib" with "include"
-    let llvm_inc_dir = llvm_lib_dir
-        .chars()
-        .take(llvm_lib_dir.len() - 3)
-        .collect::<String>()
-        + "include";
+        // replace last "lib" with "include"
+        let llvm_inc_dir = llvm_lib_dir
+            .chars()
+            .take(llvm_lib_dir.len() - 3)
+            .collect::<String>()
+            + "include";
 
-    // compile enzyme
-    let (libdir, incdir) = compile_enzyme(llvm_lib_dir.clone());
-    let libnames = [format!("EnzymeStatic-{}", llvm_version)];
+        // compile enzyme
+        let (libdir, incdir) = compile_enzyme(llvm_lib_dir.clone());
+        let libnames = [format!("EnzymeStatic-{}", llvm_version)];
 
-    // bind enzyme api
-    let bindings_rs = PathBuf::from(env::var("OUT_DIR").unwrap()).join("bindings.rs");
-    let bindings = enzyme_bindings(&[llvm_inc_dir, incdir]).expect("Couldn't generate bindings!");
-    bindings
-        .write_to_file(bindings_rs)
-        .expect("Couldn't write file bindings.rs!");
+        // bind enzyme api
+        let bindings_rs = PathBuf::from(env::var("OUT_DIR").unwrap()).join("bindings.rs");
+        let bindings =
+            enzyme_bindings(&[llvm_inc_dir, incdir]).expect("Couldn't generate bindings!");
+        bindings
+            .write_to_file(bindings_rs)
+            .expect("Couldn't write file bindings.rs!");
 
-    println!("cargo:rustc-link-search=native={}", libdir);
-    println!("cargo:rustc-link-search=native={}", llvm_lib_dir);
-    for libname in libnames.iter() {
-        println!("cargo:rustc-link-lib={}", libname);
+        println!("cargo:rustc-link-search=native={}", libdir);
+        println!("cargo:rustc-link-search=native={}", llvm_lib_dir);
+        for libname in libnames.iter() {
+            println!("cargo:rustc-link-lib={}", libname);
+        }
+        println!("cargo:rustc-link-lib=LLVMDemangle");
+        println!("cargo:rerun-if-changed=wrapper.h");
     }
-    println!("cargo:rustc-link-lib=LLVMDemangle");
-    println!("cargo:rerun-if-changed=wrapper.h");
+}
+
+fn main() {
+    #[cfg(feature = "enzyme")]
+    enzyme::enzyme_main();
 }