Merge pull request #91 from nmheim/nh/early-exit

Add parameter to disable early exit of expression evaluation

Author: MilesCranmer

benchmark/benchmarks.jl

@@ -74,11 +74,18 @@ function benchmark_evaluation()
                 extra_kws...
             )
             suite[T]["evaluation$(extra_key)"] = @benchmarkable(
-                [eval_tree_array(tree, X, $operators; turbo=$turbo, $extra_kws...) for tree in trees],
+                [eval_tree_array(tree, X, $operators; kws...) for tree in trees],
                 setup=(
                     X=randn(MersenneTwister(0), $T, 5, $n);
                     treesize=20;
                     ntrees=100;
+                    kws=$(
+                        if @isdefined(EvalOptions)
+                            (; eval_options=EvalOptions(; turbo=turbo, extra_kws...))
+                        else
+                            (; turbo, extra_kws...)
+                        end
+                    );
                     trees=[gen_random_tree_fixed_size(treesize, $operators, 5, $T) for _ in 1:ntrees]
                 )
             )

docs/src/eval.md

@@ -6,14 +6,18 @@ Given an expression tree specified with a `Node` type, you may evaluate the expr
 over an array of data with the following command:
 
 ```@docs
-eval_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum) where {T<:Number}
+eval_tree_array(
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    operators::OperatorEnum;
+    eval_options::Union{EvalOptions,Nothing}=nothing,
+) where {T}
 ```
 
-Assuming you are only using a single `OperatorEnum`, you can also use
-the following shorthand by using the expression as a function:
+You can also use the following shorthand by using the expression as a function:
 
 ```
-    (tree::AbstractExpressionNode)(X::AbstractMatrix{T}, operators::OperatorEnum; turbo::Union{Bool,Val}=false, bumper::Union{Bool,Val}=Val(false))
+    (tree::AbstractExpressionNode)(X, operators::OperatorEnum; kws...)
 
 Evaluate a binary tree (equation) over a given input data matrix. The
 operators contain all of the operators used. This function fuses doublets
@@ -23,8 +27,7 @@ and triplets of operations for lower memory usage.
 - `tree::AbstractExpressionNode`: The root node of the tree to evaluate.
 - `cX::AbstractMatrix{T}`: The input data to evaluate the tree on.
 - `operators::OperatorEnum`: The operators used in the tree.
-- `turbo::Union{Bool,Val}`: Use LoopVectorization.jl for faster evaluation.
-- `bumper::Union{Bool,Val}`: Use Bumper.jl for faster evaluation.
+- `kws...`: Passed to [`eval_tree_array`](@ref).
 
 # Returns
 - `output::AbstractVector{T}`: the result, which is a 1D array.
@@ -53,6 +56,14 @@ It also re-defines `print`, `show`, and the various operators, to work with the
     Thus, if you define an expression with one `OperatorEnum`, and then try to
     evaluate it or print it with a different `OperatorEnum`, you will get undefined behavior!
 
+    For safer behavior, you should use [`Expression`](@ref) objects.
+
+Evaluation options are specified using `EvalOptions`:
+
+```@docs
+EvalOptions
+```
+
 You can also work with arbitrary types, by defining a `GenericOperatorEnum` instead.
 The notation is the same for `eval_tree_array`, though it will return `nothing`
 when it can't find a method, and not do any NaN checks:

ext/DynamicExpressionsBumperExt.jl

@@ -2,7 +2,7 @@ module DynamicExpressionsBumperExt
 
 using Bumper: @no_escape, @alloc
 using DynamicExpressions:
-    OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array
+    OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array, EvalOptions
 using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
 import DynamicExpressions.ExtensionInterfaceModule:
@@ -12,8 +12,8 @@ function bumper_eval_tree_array(
     tree::AbstractExpressionNode{T},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
-    ::Val{turbo},
-) where {T,turbo}
+    eval_options::EvalOptions{turbo,true,early_exit},
+) where {T,turbo,early_exit}
     result = similar(cX, axes(cX, 2))
     n = size(cX, 2)
     all_ok = Ref(false)
@@ -26,7 +26,7 @@ function bumper_eval_tree_array(
                 ok = if leaf_node.constant
                     v = leaf_node.val
                     ar .= v
-                    isfinite(v)
+                    early_exit ? isfinite(v) : true
                 else
                     ar .= view(cX, leaf_node.feature, :)
                     true
@@ -38,7 +38,7 @@ function bumper_eval_tree_array(
             # In the evaluation kernel, we combine the branch nodes
             # with the arrays created by the leaf nodes:
             ((args::Vararg{Any,M}) where {M}) ->
-                dispatch_kerns!(operators, args..., Val(turbo)),
+                dispatch_kerns!(operators, args..., eval_options),
             tree;
             break_sharing=Val(true),
         )
@@ -49,55 +49,61 @@ function bumper_eval_tree_array(
     return (result, all_ok[])
 end
 
-function dispatch_kerns!(operators, branch_node, cumulator, ::Val{turbo}) where {turbo}
+function dispatch_kerns!(
+    operators, branch_node, cumulator, eval_options::EvalOptions{<:Any,true,early_exit}
+) where {early_exit}
     cumulator.ok || return cumulator
 
-    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, Val(turbo))
-    return ResultOk(out, is_valid_array(out))
+    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, eval_options)
+    return ResultOk(out, early_exit ? is_valid_array(out) : true)
 end
 function dispatch_kerns!(
-    operators, branch_node, cumulator1, cumulator2, ::Val{turbo}
-) where {turbo}
+    operators,
+    branch_node,
+    cumulator1,
+    cumulator2,
+    eval_options::EvalOptions{<:Any,true,early_exit},
+) where {early_exit}
     cumulator1.ok || return cumulator1
     cumulator2.ok || return cumulator2
 
     out = dispatch_kern2!(
-        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, Val(turbo)
+        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, eval_options
     )
-    return ResultOk(out, is_valid_array(out))
+    return ResultOk(out, early_exit ? is_valid_array(out) : true)
 end
 
-@generated function dispatch_kern1!(unaops, op_idx, cumulator, ::Val{turbo}) where {turbo}
+@generated function dispatch_kern1!(unaops, op_idx, cumulator, eval_options::EvalOptions)
     nuna = counttuple(unaops)
     quote
         Base.@nif(
             $nuna,
             i -> i == op_idx,
             i -> let op = unaops[i]
-                return bumper_kern1!(op, cumulator, Val(turbo))
+                return bumper_kern1!(op, cumulator, eval_options)
             end,
         )
     end
 end
 @generated function dispatch_kern2!(
-    binops, op_idx, cumulator1, cumulator2, ::Val{turbo}
-) where {turbo}
+    binops, op_idx, cumulator1, cumulator2, eval_options::EvalOptions
+)
     nbin = counttuple(binops)
     quote
         Base.@nif(
             $nbin,
             i -> i == op_idx,
             i -> let op = binops[i]
-                return bumper_kern2!(op, cumulator1, cumulator2, Val(turbo))
+                return bumper_kern2!(op, cumulator1, cumulator2, eval_options)
             end,
         )
     end
 end
-function bumper_kern1!(op::F, cumulator, ::Val{false}) where {F}
+function bumper_kern1!(op::F, cumulator, ::EvalOptions{false,true}) where {F}
     @. cumulator = op(cumulator)
     return cumulator
 end
-function bumper_kern2!(op::F, cumulator1, cumulator2, ::Val{false}) where {F}
+function bumper_kern2!(op::F, cumulator1, cumulator2, ::EvalOptions{false,true}) where {F}
     @. cumulator1 = op(cumulator1, cumulator2)
     return cumulator1
 end

ext/DynamicExpressionsLoopVectorizationExt.jl

@@ -3,7 +3,7 @@ module DynamicExpressionsLoopVectorizationExt
 using LoopVectorization: @turbo
 using DynamicExpressions: AbstractExpressionNode
 using DynamicExpressions.UtilsModule: ResultOk, fill_similar
-using DynamicExpressions.EvaluateModule: @return_on_check
+using DynamicExpressions.EvaluateModule: @return_on_nonfinite_val, EvalOptions
 import DynamicExpressions.EvaluateModule:
     deg1_eval,
     deg2_eval,
@@ -18,7 +18,10 @@ import DynamicExpressions.ExtensionInterfaceModule:
 _is_loopvectorization_loaded(::Int) = true
 
 function deg2_eval(
-    cumulator_l::AbstractVector{T}, cumulator_r::AbstractVector{T}, op::F, ::Val{true}
+    cumulator_l::AbstractVector{T},
+    cumulator_r::AbstractVector{T},
+    op::F,
+    ::EvalOptions{true},
 )::ResultOk where {T<:Number,F}
     @turbo for j in eachindex(cumulator_l)
         x = op(cumulator_l[j], cumulator_r[j])
@@ -28,7 +31,7 @@ function deg2_eval(
 end
 
 function deg1_eval(
-    cumulator::AbstractVector{T}, op::F, ::Val{true}
+    cumulator::AbstractVector{T}, op::F, ::EvalOptions{true}
 )::ResultOk where {T<:Number,F}
     @turbo for j in eachindex(cumulator)
         x = op(cumulator[j])
@@ -38,21 +41,25 @@ function deg1_eval(
 end
 
 function deg1_l2_ll0_lr0_eval(
-    tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{true}
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    op::F,
+    op_l::F2,
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F,F2}
     if tree.l.l.constant && tree.l.r.constant
         val_ll = tree.l.l.val
         val_lr = tree.l.r.val
-        @return_on_check val_ll cX
-        @return_on_check val_lr cX
+        @return_on_nonfinite_val(eval_options, val_ll, cX)
+        @return_on_nonfinite_val(eval_options, val_lr, cX)
         x_l = op_l(val_ll, val_lr)::T
-        @return_on_check x_l cX
+        @return_on_nonfinite_val(eval_options, x_l, cX)
         x = op(x_l)::T
-        @return_on_check x cX
+        @return_on_nonfinite_val(eval_options, x, cX)
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     elseif tree.l.l.constant
         val_ll = tree.l.l.val
-        @return_on_check val_ll cX
+        @return_on_nonfinite_val(eval_options, val_ll, cX)
         feature_lr = tree.l.r.feature
         cumulator = similar(cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
@@ -64,7 +71,7 @@ function deg1_l2_ll0_lr0_eval(
     elseif tree.l.r.constant
         feature_ll = tree.l.l.feature
         val_lr = tree.l.r.val
-        @return_on_check val_lr cX
+        @return_on_nonfinite_val(eval_options, val_lr, cX)
         cumulator = similar(cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j], val_lr)
@@ -86,15 +93,19 @@ function deg1_l2_ll0_lr0_eval(
 end
 
 function deg1_l1_ll0_eval(
-    tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{true}
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    op::F,
+    op_l::F2,
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F,F2}
     if tree.l.l.constant
         val_ll = tree.l.l.val
-        @return_on_check val_ll cX
+        @return_on_nonfinite_val(eval_options, val_ll, cX)
         x_l = op_l(val_ll)::T
-        @return_on_check x_l cX
+        @return_on_nonfinite_val(eval_options, x_l, cX)
         x = op(x_l)::T
-        @return_on_check x cX
+        @return_on_nonfinite_val(eval_options, x, cX)
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     else
         feature_ll = tree.l.l.feature
@@ -109,20 +120,23 @@ function deg1_l1_ll0_eval(
 end
 
 function deg2_l0_r0_eval(
-    tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, ::Val{true}
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    op::F,
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F}
     if tree.l.constant && tree.r.constant
         val_l = tree.l.val
-        @return_on_check val_l cX
+        @return_on_nonfinite_val(eval_options, val_l, cX)
         val_r = tree.r.val
-        @return_on_check val_r cX
+        @return_on_nonfinite_val(eval_options, val_r, cX)
         x = op(val_l, val_r)::T
-        @return_on_check x cX
+        @return_on_nonfinite_val(eval_options, x, cX)
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     elseif tree.l.constant
         cumulator = similar(cX, axes(cX, 2))
         val_l = tree.l.val
-        @return_on_check val_l cX
+        @return_on_nonfinite_val(eval_options, val_l, cX)
         feature_r = tree.r.feature
         @turbo for j in axes(cX, 2)
             x = op(val_l, cX[feature_r, j])
@@ -133,7 +147,7 @@ function deg2_l0_r0_eval(
         cumulator = similar(cX, axes(cX, 2))
         feature_l = tree.l.feature
         val_r = tree.r.val
-        @return_on_check val_r cX
+        @return_on_nonfinite_val(eval_options, val_r, cX)
         @turbo for j in axes(cX, 2)
             x = op(cX[feature_l, j], val_r)
             cumulator[j] = x
@@ -157,11 +171,11 @@ function deg2_l0_eval(
     cumulator::AbstractVector{T},
     cX::AbstractArray{T},
     op::F,
-    ::Val{true},
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F}
     if tree.l.constant
         val = tree.l.val
-        @return_on_check val cX
+        @return_on_nonfinite_val(eval_options, val, cX)
         @turbo for j in eachindex(cumulator)
             x = op(val, cumulator[j])
             cumulator[j] = x
@@ -182,11 +196,11 @@ function deg2_r0_eval(
     cumulator::AbstractVector{T},
     cX::AbstractArray{T},
     op::F,
-    ::Val{true},
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F}
     if tree.r.constant
         val = tree.r.val
-        @return_on_check val cX
+        @return_on_nonfinite_val(eval_options, val, cX)
         @turbo for j in eachindex(cumulator)
             x = op(cumulator[j], val)
             cumulator[j] = x
@@ -203,11 +217,15 @@ function deg2_r0_eval(
 end
 
 ## Interface with Bumper.jl
-function bumper_kern1!(op::F, cumulator, ::Val{true}) where {F}
+function bumper_kern1!(
+    op::F, cumulator, ::EvalOptions{true,true,early_exit}
+) where {F,early_exit}
     @turbo @. cumulator = op(cumulator)
     return cumulator
 end
-function bumper_kern2!(op::F, cumulator1, cumulator2, ::Val{true}) where {F}
+function bumper_kern2!(
+    op::F, cumulator1, cumulator2, ::EvalOptions{true,true,early_exit}
+) where {F,early_exit}
     @turbo @. cumulator1 = op(cumulator1, cumulator2)
     return cumulator1
 end

src/DynamicExpressions.jl

@@ -70,7 +70,8 @@ import .NodeModule:
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!
-@reexport import .EvaluateModule: eval_tree_array, differentiable_eval_tree_array
+@reexport import .EvaluateModule:
+    eval_tree_array, differentiable_eval_tree_array, EvalOptions
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
 @reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!