ArenaNode for array-backed expressions

Project.toml

@@ -41,7 +41,7 @@ Optim = "1, 2"
 NLSolversBase = "7, 8"
 PrecompileTools = "1.2.1"
 Reexport = "1.2.2"
-SymbolicUtils = "4"
+SymbolicUtils = "4.35"
 Zygote = "0.7"
 julia = "1.10"
 Random = "1"

ext/DynamicExpressionsLoopVectorizationExt.jl

@@ -211,12 +211,16 @@ function deg2_r0_eval(
 end
 
 # Interface with Bumper.jl
-function bumper_kern!(
+@generated function bumper_kern!(
     op::F, cumulators::Tuple{Vararg{Any,degree}}, ::EvalOptions{true,true,early_exit}
 ) where {F,degree,early_exit}
-    cumulator_1 = first(cumulators)
-    @turbo @. cumulator_1 = op(cumulators...)
-    return cumulator_1
+    quote
+        Base.Cartesian.@nexprs($degree, i -> cumulator_i = cumulators[i])
+        @turbo for j in eachindex(cumulator_1)
+            cumulator_1[j] = Base.Cartesian.@ncall($degree, op, i -> cumulator_i[j])
+        end
+        return cumulator_1
+    end
 end
 
 end

scripts/bench_arenanode_eval.jl

@@ -0,0 +1,89 @@
+using DynamicExpressions
+using DynamicExpressions: EvalOptions, ArrayBuffer
+using Random: MersenneTwister
+
+include(joinpath(@__DIR__, "..", "test", "tree_gen_utils.jl"))
+
+const AN = DynamicExpressions.ArenaNodeModule
+
+operators = OperatorEnum(1 => (cos, exp), 2 => (+, -, *, /))
+const T = Float64
+nfeat = 5
+n = 1_000
+rng = MersenneTwister(0)
+X = randn(rng, T, nfeat, n)
+buf = zeros(T, 64, n)
+
+function bench(trees, X, operators, buffer; reps=300)
+    best = Inf
+    for _ in 1:reps
+        t0 = time_ns()
+        for tree in trees
+            buffer.index[] = 0
+            eval_tree_array(tree, X, operators; eval_options=EvalOptions(; buffer))
+        end
+        best = min(best, (time_ns() - t0) / length(trees))
+    end
+    return best
+end
+
+function allocs_per_eval(tree, X, operators, buffer)
+    @allocated(eval_tree_array(tree, X, operators; eval_options=EvalOptions(; buffer)))
+end
+
+for treesize in (7, 15, 31)
+    trees = [gen_random_tree_fixed_size(treesize, operators, nfeat, T) for _ in 1:50]
+    atrees = [convert(AN.ArenaNode{T,2}, t) for t in trees]
+    buffer = ArrayBuffer(buf, Ref(0))
+
+    # correctness sanity: both paths must agree
+    for (t, a) in zip(trees, atrees)
+        buffer.index[] = 0
+        yn, okn = eval_tree_array(t, X, operators; eval_options=EvalOptions(; buffer))
+        buffer.index[] = 0
+        ya, oka = eval_tree_array(a, X, operators; eval_options=EvalOptions(; buffer))
+        okn == oka || error("ok mismatch at size $treesize")
+        okn && (yn ≈ ya || error("value mismatch at size $treesize"))
+    end
+
+    # warmup
+    bench(trees, X, operators, buffer; reps=3)
+    bench(atrees, X, operators, buffer; reps=3)
+
+    t_node = bench(trees, X, operators, buffer)
+    t_arena = bench(atrees, X, operators, buffer)
+    buffer.index[] = 0
+    a_node = allocs_per_eval(trees[1], X, operators, buffer)
+    buffer.index[] = 0
+    a_arena = allocs_per_eval(atrees[1], X, operators, buffer)
+    println(
+        "n=$treesize  Node: $(round(t_node/1e3; digits=2))us  " *
+        "ArenaNode: $(round(t_arena/1e3; digits=2))us  " *
+        "ratio=$(round(t_arena/t_node; digits=3))  " *
+        "allocs/eval Node=$a_node Arena=$a_arena",
+    )
+
+    bench_nobuf(trees) = begin
+        best = Inf
+        for _ in 1:300
+            t0 = time_ns()
+            for tree in trees
+                eval_tree_array(tree, X, operators)
+            end
+            best = min(best, (time_ns() - t0) / length(trees))
+        end
+        best
+    end
+    bench_nobuf(trees[1:2])
+    bench_nobuf(atrees[1:2])  # warmup
+    tn_nb = bench_nobuf(trees)
+    ta_nb = bench_nobuf(atrees)
+    an_nb = @allocated(eval_tree_array(trees[1], X, operators))
+    aa_nb = @allocated(eval_tree_array(atrees[1], X, operators))
+    println(
+        "n=$treesize  unbuffered:  Node: $(round(tn_nb/1e3; digits=2))us  " *
+        "ArenaNode: $(round(ta_nb/1e3; digits=2))us  " *
+        "ratio=$(round(ta_nb/tn_nb; digits=3))  " *
+        "allocs/eval Node=$an_nb Arena=$aa_nb",
+    )
+end