Add parameter to disable early exit of expression evaluation

ext/DynamicExpressionsBumperExt.jl

@@ -1,7 +1,7 @@
 module DynamicExpressionsBumperExt
 
 using Bumper: @no_escape, @alloc
-using DynamicExpressions: OperatorEnum, AbstractExpressionNode, tree_mapreduce
+using DynamicExpressions: OperatorEnum, AbstractExpressionNode, tree_mapreduce, EvaluationOptions
 using DynamicExpressions.UtilsModule: ResultOk, counttuple, is_bad_array
 
 import DynamicExpressions.ExtensionInterfaceModule:
@@ -11,8 +11,8 @@ function bumper_eval_tree_array(
     tree::AbstractExpressionNode{T},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
-    ::Val{turbo},
-) where {T,turbo}
+    options::EvaluationOptions{turbo,true,early_exit}
+) where {T,turbo,early_exit}
     result = similar(cX, axes(cX, 2))
     n = size(cX, 2)
     all_ok = Ref(false)
@@ -25,7 +25,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
@@ -37,7 +37,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..., options),
             tree;
             break_sharing=Val(true),
         )
@@ -48,55 +48,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, options::EvaluationOptions{turbo,true,early_exit}
+) where {turbo,early_exit}
     cumulator.ok || return cumulator
 
-    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, Val(turbo))
-    return ResultOk(out, !is_bad_array(out))
+    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, options)
+    return early_exit ? ResultOk(out, !is_bad_array(out)) : ResultOk(out, true)
 end
 function dispatch_kerns!(
-    operators, branch_node, cumulator1, cumulator2, ::Val{turbo}
-) where {turbo}
+    operators, branch_node, cumulator1, cumulator2, options::EvaluationOptions{turbo,true,early_exit}
+) where {turbo,early_exit}
     cumulator1.ok || return cumulator1
     cumulator2.ok || return cumulator2
 
-    out = dispatch_kern2!(
-        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, Val(turbo)
-    )
-    return ResultOk(out, !is_bad_array(out))
+    out = dispatch_kern2!(operators.binops, branch_node.op, cumulator1.x, cumulator2.x, options)
+    return early_exit ? ResultOk(out, !is_bad_array(out)) : ResultOk(out, true)
 end
 
-@generated function dispatch_kern1!(unaops, op_idx, cumulator, ::Val{turbo}) where {turbo}
+@generated function dispatch_kern1!(
+    unaops, op_idx, cumulator, options::EvaluationOptions{turbo,true,early_exit}
+) where {turbo,early_exit}
     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, options)
             end,
         )
     end
 end
 @generated function dispatch_kern2!(
-    binops, op_idx, cumulator1, cumulator2, ::Val{turbo}
-) where {turbo}
+    binops, op_idx, cumulator1, cumulator2, options::EvaluationOptions{turbo,true,early_exit}
+) where {turbo,early_exit}
     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, options)
             end,
         )
     end
 end
-function bumper_kern1!(op::F, cumulator, ::Val{false}) where {F}
+function bumper_kern1!(
+    op::F, cumulator, ::EvaluationOptions{false,true,early_exit}
+) where {F,early_exit}
     @. cumulator = op(cumulator)
     return cumulator
 end
-function bumper_kern2!(op::F, cumulator1, cumulator2, ::Val{false}) where {F}
+function bumper_kern2!(
+    op::F, cumulator1, cumulator2, ::EvaluationOptions{false,true,early_exit}
+) where {F,early_exit}
     @. 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_check, EvaluationOptions
 import DynamicExpressions.EvaluateModule:
     deg1_eval,
     deg2_eval,
@@ -18,7 +18,7 @@ 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, ::EvaluationOptions{true}
 )::ResultOk where {T<:Number,F}
     @turbo for j in eachindex(cumulator_l)
         x = op(cumulator_l[j], cumulator_r[j])
@@ -28,7 +28,7 @@ function deg2_eval(
 end
 
 function deg1_eval(
-    cumulator::AbstractVector{T}, op::F, ::Val{true}
+    cumulator::AbstractVector{T}, op::F, ::EvaluationOptions{true}
 )::ResultOk where {T<:Number,F}
     @turbo for j in eachindex(cumulator)
         x = op(cumulator[j])
@@ -38,7 +38,7 @@ 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, ::EvaluationOptions{true}
 ) where {T<:Number,F,F2}
     if tree.l.l.constant && tree.l.r.constant
         val_ll = tree.l.l.val
@@ -86,7 +86,7 @@ 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, ::EvaluationOptions{true}
 ) where {T<:Number,F,F2}
     if tree.l.l.constant
         val_ll = tree.l.l.val
@@ -109,7 +109,7 @@ 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, ::EvaluationOptions{true}
 ) where {T<:Number,F}
     if tree.l.constant && tree.r.constant
         val_l = tree.l.val
@@ -157,7 +157,7 @@ function deg2_l0_eval(
     cumulator::AbstractVector{T},
     cX::AbstractArray{T},
     op::F,
-    ::Val{true},
+    ::EvaluationOptions{true}
 ) where {T<:Number,F}
     if tree.l.constant
         val = tree.l.val
@@ -182,7 +182,7 @@ function deg2_r0_eval(
     cumulator::AbstractVector{T},
     cX::AbstractArray{T},
     op::F,
-    ::Val{true},
+    ::EvaluationOptions{true}
 ) where {T<:Number,F}
     if tree.r.constant
         val = tree.r.val
@@ -203,11 +203,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, ::EvaluationOptions{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, ::EvaluationOptions{true,true,early_exit}
+) where {F,early_exit}
     @turbo @. cumulator1 = op(cumulator1, cumulator2)
     return cumulator1
 end

src/DynamicExpressions.jl

@@ -59,7 +59,7 @@ 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, EvaluationOptions
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
 @reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!