feat: compat with turbo mode and buffered evals

Author: MilesCranmer

ext/DynamicExpressionsLoopVectorizationExt.jl

@@ -2,8 +2,9 @@ module DynamicExpressionsLoopVectorizationExt
 
 using LoopVectorization: @turbo
 using DynamicExpressions: AbstractExpressionNode
-using DynamicExpressions.UtilsModule: ResultOk, fill_similar
-using DynamicExpressions.EvaluateModule: @return_on_nonfinite_val, EvalOptions
+using DynamicExpressions.UtilsModule: ResultOk
+using DynamicExpressions.EvaluateModule:
+    @return_on_nonfinite_val, EvalOptions, get_array, get_feature_array, get_filled_array
 import DynamicExpressions.EvaluateModule:
     deg1_eval,
     deg2_eval,
@@ -56,12 +57,12 @@ function deg1_l2_ll0_lr0_eval(
         @return_on_nonfinite_val(eval_options, x_l, cX)
         x = op(x_l)::T
         @return_on_nonfinite_val(eval_options, x, cX)
-        return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
+        return ResultOk(get_filled_array(eval_options.buffer, x, cX, axes(cX, 2)))
     elseif tree.l.l.constant
         val_ll = tree.l.l.val
         @return_on_nonfinite_val(eval_options, val_ll, cX)
         feature_lr = tree.l.r.feature
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
             x_l = op_l(val_ll, cX[feature_lr, j])
             x = op(x_l)
@@ -72,7 +73,7 @@ function deg1_l2_ll0_lr0_eval(
         feature_ll = tree.l.l.feature
         val_lr = tree.l.r.val
         @return_on_nonfinite_val(eval_options, val_lr, cX)
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j], val_lr)
             x = op(x_l)
@@ -82,7 +83,7 @@ function deg1_l2_ll0_lr0_eval(
     else
         feature_ll = tree.l.l.feature
         feature_lr = tree.l.r.feature
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j], cX[feature_lr, j])
             x = op(x_l)
@@ -106,10 +107,10 @@ function deg1_l1_ll0_eval(
         @return_on_nonfinite_val(eval_options, x_l, cX)
         x = op(x_l)::T
         @return_on_nonfinite_val(eval_options, x, cX)
-        return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
+        return ResultOk(get_filled_array(eval_options.buffer, x, cX, axes(cX, 2)))
     else
         feature_ll = tree.l.l.feature
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j])
             x = op(x_l)
@@ -132,9 +133,9 @@ function deg2_l0_r0_eval(
         @return_on_nonfinite_val(eval_options, val_r, cX)
         x = op(val_l, val_r)::T
         @return_on_nonfinite_val(eval_options, x, cX)
-        return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
+        return ResultOk(get_filled_array(eval_options.buffer, x, cX, axes(cX, 2)))
     elseif tree.l.constant
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         val_l = tree.l.val
         @return_on_nonfinite_val(eval_options, val_l, cX)
         feature_r = tree.r.feature
@@ -144,7 +145,7 @@ function deg2_l0_r0_eval(
         end
         return ResultOk(cumulator, true)
     elseif tree.r.constant
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         feature_l = tree.l.feature
         val_r = tree.r.val
         @return_on_nonfinite_val(eval_options, val_r, cX)
@@ -154,7 +155,7 @@ function deg2_l0_r0_eval(
         end
         return ResultOk(cumulator, true)
     else
-        cumulator = similar(cX, axes(cX, 2))
+        cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         feature_l = tree.l.feature
         feature_r = tree.r.feature
         @turbo for j in axes(cX, 2)

src/Evaluate.jl

@@ -106,7 +106,7 @@ end
     v_bumper = _to_bool_val(bumper)
     v_early_exit = _to_bool_val(early_exit)
 
-    if v_turbo isa Val{true} || v_bumper isa Val{true}
+    if v_bumper isa Val{true}
         @assert buffer === nothing && buffer_ref === nothing
     end
 

src/EvaluateDerivative.jl

@@ -5,7 +5,8 @@ import ..OperatorEnumModule: OperatorEnum
 import ..UtilsModule: fill_similar, ResultOk2
 import ..ValueInterfaceModule: is_valid_array
 import ..NodeUtilsModule: count_constant_nodes, index_constant_nodes, NodeIndex
-import ..EvaluateModule: deg0_eval, get_nuna, get_nbin, OPERATOR_LIMIT_BEFORE_SLOWDOWN
+import ..EvaluateModule:
+    deg0_eval, get_nuna, get_nbin, OPERATOR_LIMIT_BEFORE_SLOWDOWN, EvalOptions
 import ..ExtensionInterfaceModule: _zygote_gradient
 
 """
@@ -120,7 +121,7 @@ end
 function diff_deg0_eval(
     tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, direction::Integer
 ) where {T<:Number}
-    const_part = deg0_eval(tree, cX).x
+    const_part = deg0_eval(tree, cX, EvalOptions()).x
     derivative_part = if ((!tree.constant) && tree.feature == direction)
         fill_similar(one(T), cX, axes(cX, 2))
     else
@@ -335,7 +336,7 @@ function grad_deg0_eval(
     cX::AbstractMatrix{T},
     ::Val{mode},
 )::ResultOk2 where {T<:Number,mode}
-    const_part = deg0_eval(tree, cX).x
+    const_part = deg0_eval(tree, cX, EvalOptions()).x
 
     zero_mat = if isa(cX, Array)
         fill_similar(zero(T), cX, n_gradients, axes(cX, 2))

test/test_buffered_evaluation.jl

@@ -119,6 +119,7 @@ end
 @testitem "Random tree buffer evaluation" begin
     using DynamicExpressions
     using Random
+    using LoopVectorization
     include("tree_gen_utils.jl")
 
     # Test setup
@@ -127,18 +128,19 @@ end
         binary_operators=[+, -, *, /], unary_operators=[sin, cos, exp]
     )
 
-    for i in 1:100
+    for turbo in (false, true), i in 1:100
         # Generate a random tree with varying size (1-10 nodes)
         n_nodes = rand(1:10)
         tree = gen_random_tree_fixed_size(n_nodes, operators, size(X, 1), Float64, Node)
 
         # Regular evaluation
-        result1, ok1 = eval_tree_array(tree, X, operators)
+        eval_options_no_buffer = EvalOptions(; turbo)
+        result1, ok1 = eval_tree_array(tree, X, operators; eval_options=eval_options_no_buffer)
 
         # Buffer evaluation
         buffer = Array{Float64}(undef, 2n_nodes, size(X, 2))
         buffer_ref = Ref(rand(1:10))  # Random starting index (will be reset)
-        eval_options = EvalOptions(; buffer=buffer, buffer_ref=buffer_ref)
+        eval_options = EvalOptions(; turbo, buffer, buffer_ref)
         result2, ok2 = eval_tree_array(tree, X, operators; eval_options)
 
         # Results should be identical