Merge pull request #12 from martinjrobins/bench

add optimisations and benchmarking

Cargo.toml

@@ -69,6 +69,13 @@ llvm-sys-170 = { package = "llvm-sys", version = "170.0.1", optional = true }
 [build-dependencies]
 cmake = "0.1.50"
 
+[dev-dependencies]
+divan = "0.1.14"
+
+[[bench]]
+name = "evaluation"
+harness = false
+
 [profile.dev]
 opt-level = 0
 debug = true

benches/evaluation.rs

@@ -0,0 +1,58 @@
+use diffsl::{discretise::DiscreteModel, execution::Compiler, parser::parse_ds_string};
+use divan::Bencher;
+use ndarray::Array1;
+
+fn main() {
+    divan::main();
+}
+
+fn setup(n: usize, f_text: &str, name: &str) -> Compiler {
+    let u = vec![1.0; n];
+    let full_text = format!(
+        "
+        u_i {{
+            {} 
+        }}
+        F_i {{
+            {}
+        }}
+        out_i {{
+            u_i 
+        }}
+        ",
+        (0..n)
+            .map(|i| format!("x{} = {},", i, u[i]))
+            .collect::<Vec<_>>()
+            .join("\n"),
+        f_text,
+    );
+    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()
+}
+
+#[divan::bench(consts = [1, 10, 100, 1000])]
+fn add_scalar_diffsl<const N: usize>(bencher: Bencher) {
+    let n = N;
+    let compiler = setup(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_ndarray<const N: usize>(bencher: Bencher) {
+    let n = N;
+    let u = Array1::from_shape_vec((n,), vec![1.0; n]).unwrap();
+
+    bencher.bench_local(|| {
+        let _ = &u + 1.0;
+    });
+}

src/ast/mod.rs

@@ -435,12 +435,24 @@ impl<'a> AstKind<'a> {
             args: vec![Box::new(child)],
         })
     }
+    pub fn new_indice(first: Ast<'a>, last: Option<Ast<'a>>, sep: Option<&'a str>) -> Self {
+        AstKind::Indice(Indice {
+            first: Box::new(first),
+            last: last.map(Box::new),
+            sep,
+        })
+    }
     pub fn new_index(left: Ast<'a>, right: Ast<'a>) -> Self {
         AstKind::Index(Index {
             left: Box::new(left),
             right: Box::new(right),
         })
     }
+    pub fn new_vector(data: Vec<Ast<'a>>) -> Self {
+        AstKind::Vector(Vector {
+            data: data.into_iter().map(Box::new).collect(),
+        })
+    }
     pub fn new_name(name: &'a str) -> Self {
         AstKind::Name(name)
     }
@@ -471,6 +483,9 @@ impl<'a> AstKind<'a> {
     pub fn new_num(num: f64) -> Self {
         AstKind::Number(num)
     }
+    pub fn new_integer(num: i64) -> Self {
+        AstKind::Integer(num)
+    }
     pub fn new_tensor(name: &'a str, indices: Vec<char>, elmts: Vec<Ast<'a>>) -> Self {
         AstKind::Tensor(Tensor::new(name, indices, elmts))
     }

src/discretise/discrete_model.rs

@@ -408,8 +408,29 @@ impl<'s> DiscreteModel<'s> {
                 kind: AstKind::new_num(0.0),
                 span: None,
             };
+            let indices = if ret.state.shape().is_empty() {
+                None
+            } else {
+                let first = Ast {
+                    kind: AstKind::new_integer(0),
+                    span: None,
+                };
+                let last = Ast {
+                    kind: AstKind::new_integer(i64::try_from(ret.state.shape()[0]).unwrap()),
+                    span: None,
+                };
+                let index = Ast {
+                    kind: AstKind::new_indice(first, Some(last), Some(":")),
+                    span: None,
+                };
+                Some(Ast {
+                    kind: AstKind::new_vector(vec![index]),
+                    span: None,
+                })
+            };
+
             let zero_block = Ast {
-                kind: AstKind::new_tensor_elmt(zero, None),
+                kind: AstKind::new_tensor_elmt(zero, indices),
                 span: None,
             };
             let indices = ret.state.indices().to_vec();
@@ -966,7 +987,7 @@ mod tests {
                 y,
             }
         " [],
-        logistic_no_f: "
+        logistic_no_m: "
             in = [r, ]
             r { 1, }    
             u {
@@ -979,13 +1000,15 @@ mod tests {
                 y,
             }
         " [],
-        logistic_no_f2: "
+        logistic_no_m2: "
             in = [r, ]
             r { 1, }    
             u_i {
+                x = 1,
                 y = 1,
             }
             F_i {
+                (r * x) * (1 - x),
                 (r * y) * (1 - y),
             }
             out {

src/execution/compiler.rs

@@ -193,6 +193,7 @@ impl Compiler {
                     .arg(format!("-load={}", enzyme_lib))
                     .arg("-enzyme")
                     .arg("--enable-new-pm=0")
+                    .arg("-O3")
                     .arg("-o")
                     .arg(bitcodefilename.as_str())
                     .output()?;
@@ -209,7 +210,7 @@ impl Compiler {
                 let module = Module::parse_bitcode_from_buffer(&buffer, context)
                     .map_err(|e| anyhow::anyhow!("Error parsing bitcode: {:?}", e))?;
                 let ee = module
-                    .create_jit_execution_engine(OptimizationLevel::None)
+                    .create_jit_execution_engine(OptimizationLevel::Default)
                     .map_err(|e| anyhow::anyhow!("Error creating execution engine: {:?}", e))?;
 
                 let set_u0 = Compiler::jit("set_u0", &ee)?;