diff --git a/Cargo.toml b/Cargo.toml index 7495c02..c0a2516 100644 --- a/Cargo.toml +++ b/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 diff --git a/benches/evaluation.rs b/benches/evaluation.rs new file mode 100644 index 0000000..5546228 --- /dev/null +++ b/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::>() + .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(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(bencher: Bencher) { + let n = N; + let u = Array1::from_shape_vec((n,), vec![1.0; n]).unwrap(); + + bencher.bench_local(|| { + let _ = &u + 1.0; + }); +} diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 3981107..e134806 100644 --- a/src/ast/mod.rs +++ b/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>, 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>) -> 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, elmts: Vec>) -> Self { AstKind::Tensor(Tensor::new(name, indices, elmts)) } diff --git a/src/discretise/discrete_model.rs b/src/discretise/discrete_model.rs index d3bc278..dc12984 100644 --- a/src/discretise/discrete_model.rs +++ b/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 { diff --git a/src/execution/compiler.rs b/src/execution/compiler.rs index 1f36084..7a02e3d 100644 --- a/src/execution/compiler.rs +++ b/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)?;