fix: use eval instead of RGFs

Author: AayushSabharwal

benchmarks/Symbolics/BCR.jmd

@@ -33,7 +33,7 @@ rn    = complete(prnbng.rn; split = false)
 obs = [eq.lhs for eq in observed(rn)]
 osys = convert(ODESystem, rn)
 
-rhs = [eq.rhs for eq in full_equations(osys)]
+rhs = [eq.rhs for eq in full_equations(osys)][1:10]
 vars = unknowns(osys)
 pars = parameters(osys)
 
@@ -55,10 +55,13 @@ args = (vars, pars, ModelingToolkit.get_iv(osys))
 # out of place versions run into an error saying the expression is too large
 # due to the `SymbolicUtils.Code.create_array` call. `iip_config` prevents it
 # from trying to build the function.
-kwargs = (; iip_config = (false, true), expression = Val{false})
+kwargs = (; iip_config = (false, true), expression = Val{true})
 @timeit to "Build jacobian - no CSE" _, jac_nocse_iip = build_function(jac, args...; cse = false, kwargs...);
 @timeit to "Build jacobian - CSE" _, jac_cse_iip = build_function(jac, args...; cse = true, kwargs...);
 
+jac_nocse_iip = eval(jac_nocse_iip)
+jac_cse_iip = eval(jac_cse_iip)
+
 defs = defaults(osys)
 u = Float64[Symbolics.fixpoint_sub(var, defs) for var in vars]
 buffer_cse = similar(jac, Float64)
@@ -81,7 +84,7 @@ We'll also measure scaling.
 
 
 ```julia
-function run_and_time!(rhs, vars, pars, iv, u0, p, t0, N, i, jac_times, jac_allocs, build_times, first_call_times, second_call_times)
+function run_and_time_construct!(rhs, vars, pars, iv, N, i, jac_times, jac_allocs, build_times, functions)
     outputs = rhs[1:N]
     SymbolicUtils.ENABLE_HASHCONSING[] = false
     jac_result = @be Symbolics.sparsejacobian(outputs, vars)
@@ -99,46 +102,70 @@ function run_and_time!(rhs, vars, pars, iv, u0, p, t0, N, i, jac_times, jac_allo
     @assert isequal(jac_nohc, jac_hc)
     jac = jac_hc
     args = (vars, pars, iv)
-    kwargs = (; iip_config = (false, true), expression = Val{false})
+    kwargs = (; iip_config = (false, true), expression = Val{true})
 
     build_result = @be build_function(jac, args...; cse = false, kwargs...);
     build_times[1][i] = minimum(x -> x.time, build_result.samples)
-    jacfn_nocse = build_function(jac, args...; cse = false, kwargs...)[2]
+    jacfn_nocse = eval(build_function(jac, args...; cse = false, kwargs...)[2])
 
     build_result = @be build_function(jac, args...; cse = true, kwargs...);
     build_times[2][i] = minimum(x -> x.time, build_result.samples)
-    jacfn_cse = build_function(jac, args...; cse = true, kwargs...)[2]
+    jacfn_cse = eval(build_function(jac, args...; cse = true, kwargs...)[2])
+
+    functions[1][i] = let buffer = similar(jac, Float64), fn = jacfn_nocse
+        function nocse(u, p, t)
+            fn(buffer, u, p, t)
+            buffer
+        end
+    end
+    functions[2][i] = let buffer = similar(jac, Float64), fn = jacfn_cse
+        function cse(u, p, t)
+            fn(buffer, u, p, t)
+            buffer
+        end
+    end
+
+    return nothing
+end
+
+function run_and_time_call!(i, u, p, tt, functions, first_call_times, second_call_times)
+    jacfn_nocse = functions[1][i]
+    jacfn_cse = functions[2][i]
 
-    buffer = similar(jac, Float64)
-    call_result = @timed jacfn_nocse(buffer, u0, p, t0)
+    call_result = @timed jacfn_nocse(u, p, tt)
     first_call_times[1][i] = call_result.time
-    call_result = @timed jacfn_cse(buffer, u0, p, t0)
+    call_result = @timed jacfn_cse(u, p, tt)
     first_call_times[2][i] = call_result.time
 
-    call_result = @be jacfn_nocse(buffer, u0, p, t0)
+    call_result = @be jacfn_nocse(u, p, tt)
     second_call_times[1][i] = minimum(x -> x.time, call_result.samples)
-    call_result = @be jacfn_cse(buffer, u0, p, t0)
+    call_result = @be jacfn_cse(u, p, tt)
     second_call_times[2][i] = minimum(x -> x.time, call_result.samples)
-
-    return nothing
 end
 ```
 
 # Run benchmark
 
 ```julia
-N = [10, 20, 40, 60, 100, 200, 300, 400]
+N = [10, 20, 40, 80]
 jacobian_times = [zeros(Float64, length(N)), zeros(Float64, length(N))]
 jacobian_allocs = copy.(jacobian_times)
+functions = [Vector{Any}(undef, length(N)), Vector{Any}(undef, length(N))]
 # [without_cse_times, with_cse_times]
 build_times = copy.(jacobian_times)
 first_call_times = copy.(jacobian_times)
 second_call_times = copy.(jacobian_times)
 
 iv = ModelingToolkit.get_iv(osys)
-run_and_time!(rhs, vars, pars, iv, u, p, tt, 10, 1, jacobian_times, jacobian_allocs, build_times, first_call_times, second_call_times)
+run_and_time_construct!(rhs, vars, pars, iv, 10, 1, jacobian_times, jacobian_allocs, build_times, functions)
+run_and_time_call!(1, u, p, tt, functions, first_call_times, second_call_times)
+for (i, n) in enumerate(N)
+    @info i n
+    run_and_time_construct!(rhs, vars, pars, iv, n, i, jacobian_times, jacobian_allocs, build_times, functions)
+end
 for (i, n) in enumerate(N)
-    run_and_time!(rhs, vars, pars, iv, u, p, tt, n, i, jacobian_times, jacobian_allocs, build_times, first_call_times, second_call_times)
+    @info i n
+    run_and_time_call!(i, u, p, tt, functions, first_call_times, second_call_times)
 end
 ```
 

benchmarks/Symbolics/ThermalFluid.jmd

@@ -251,9 +251,9 @@ function call(fn, args...)
   fn(args...)
 end
 
-function run_and_time!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_times, first_call_times, second_call_times, i, N)
+function run_and_time_construction!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_times, functions, i, N)
   @mtkbuild sys = TestBenchPreinsulated(L=470, N=N, dn=0.3127, t_layer=[0.0056, 0.058])
-  @info "Built system" N
+  @info "Built system"
   SymbolicUtils.ENABLE_HASHCONSING[] = false
   jac_result = @be calculate_jacobian(sys; sparse = true)
   @info "No hashconsing benchmark"
@@ -287,39 +287,58 @@ function run_and_time!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_
   buffer_cse = similar(jac, Float64)
   buffer_cse.nzval .= 0.0
 
-  f_jac_nocse = build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{false}, cse = false)[2]
+  f_jac_nocse = eval(build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{true}, cse = false)[2])
+  functions[1][i] = let buffer_nocse = buffer_nocse, u0 = u0, p = p, t0 = t0, f_jac_nocse = f_jac_nocse
+    function nocse()
+      f_jac_nocse(buffer_nocse, u0, p, t0)
+      buffer_nocse
+    end
+  end
   @info "No CSE build_function result"
-  build_result_nocse = @be build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{false}, cse = false)
+  build_result_nocse = @be build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{true}, cse = false)
   @info "No CSE build_function benchmark"
   build_times[1][i] = mean(x -> x.time, build_result_nocse.samples)
-  @info "time" build_times[1][i]
-  first_call_result_nocse = @timed call(f_jac_nocse, buffer_nocse, u0, p, t0)
-  first_call_times[1][i] = first_call_result_nocse.time
-  @info "time" first_call_times[1][i]
-  second_call_result_nocse = @be call(f_jac_nocse, buffer_nocse, u0, p, t0)
-  second_call_times[1][i] = mean(x -> x.time, second_call_result_nocse.samples)
-  @info "time" second_call_times[1][i]
-
-  f_jac_cse = build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{false}, cse = true)[2]
+  @info "build_function time" build_times[1][i]
+
+  f_jac_cse = eval(build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{true}, cse = true)[2])
+  functions[2][i] = let buffer_cse = buffer_cse, u0 = u0, p = p, t0 = t0, f_jac_cse = f_jac_cse
+    function nocse()
+      f_jac_cse(buffer_cse, u0, p, t0)
+      buffer_cse
+    end
+  end
   @info "CSE build_function result"
-  build_result_cse = @be build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{false}, cse = true)
+  build_result_cse = @be build_function(jac, dvs, ps, t; iip_config = (false, true), expression = Val{true}, cse = true)
   @info "CSE build_function benchmark"
   build_times[2][i] = mean(x -> x.time, build_result_cse.samples)
-  @info "time" build_times[2][i]
-  first_call_result_cse = @timed call(f_jac_cse, buffer_cse, u0, p, t0)
-  first_call_times[2][i] = first_call_result_cse.time
-  @info "time" first_call_times[2][i]
-  second_call_result_cse = @be call(f_jac_cse, buffer_cse, u0, p, t0)
-  second_call_times[2][i] = mean(x -> x.time, second_call_result_cse.samples)
-  @info "time" second_call_times[2][i]
+  @info "build_function time" build_times[2][i]
 
   return nothing
 end
+
+function run_and_time_call!(functions, first_call_times, second_call_times, i)
+  fnocse = functions[1][i]
+  fcse = functions[2][i]
+  first_call_result_nocse = @timed fnocse()
+  first_call_times[1][i] = first_call_result_nocse.time
+  @info "First call time" first_call_times[1][i]
+  second_call_result_nocse = @be fnocse()
+  second_call_times[1][i] = mean(x -> x.time, second_call_result_nocse.samples)
+  @info "Runtime" second_call_times[1][i]
+
+  first_call_result_cse = @timed fcse()
+  first_call_times[2][i] = first_call_result_cse.time
+  @info "First call time" first_call_times[2][i]
+  second_call_result_cse = @be fcse()
+  second_call_times[2][i] = mean(x -> x.time, second_call_result_cse.samples)
+  @info "Runtime" second_call_times[2][i]
+end
 ```
 
 ```julia
-N = [5, 10, 20, 40, 80, 160, 320];
+N = [5, 10, 20, 40];
 jacobian_times = [zeros(Float64, length(N)), zeros(Float64, length(N))]
+functions = [Vector{Any}(undef, length(N)), Vector{Any}(undef, length(N))]
 jacobian_gctimes = copy.(jacobian_times)
 jacobian_allocs = copy.(jacobian_times)
 # [without_cse_times, with_cse_times]
@@ -332,10 +351,16 @@ second_call_times = copy.(jacobian_times)
 
 ```julia
 # compile
-run_and_time!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_times, first_call_times, second_call_times, 1, 5)
+run_and_time_construction!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_times, functions, 1, 5)
+run_and_time_call!(functions, first_call_times, second_call_times, 1)
+for (i, n) in enumerate(N)
+  @info i n
+  @time run_and_time_construction!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_times, functions, i, n)
+end
+
 for (i, n) in enumerate(N)
   @info i n
-  @time run_and_time!(jacobian_times, jacobian_gctimes, jacobian_allocs, build_times, first_call_times, second_call_times, i, n)
+  run_and_time_call!(functions, first_call_times, second_call_times, i)
 end
 ```