Merge pull request #112 from SymbolicML/buffered-evals

Zero-allocation tree evaluation with buffer

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)), true)
     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)), true)
     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)), true)
     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/DynamicExpressions.jl

@@ -73,6 +73,7 @@ import .StringsModule: get_op_name
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!
 @reexport import .EvaluateModule:
     eval_tree_array, differentiable_eval_tree_array, EvalOptions
+import .EvaluateModule: ArrayBuffer
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
 @reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!