Merge pull request #10 from SymbolicML/cleaner-barriers

Reduce specialization in evaluation methods

Author: MilesCranmer

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/DynamicExpressions.jl

@@ -39,4 +39,8 @@ const PACKAGE_VERSION = let
     VersionNumber(project["version"])
 end
 
+macro ignore(args...) end
+# To get LanguageServer to register library within tests
+@ignore include("../test/runtests.jl")
+
 end

src/EvaluateEquation.jl

@@ -3,7 +3,7 @@ module EvaluateEquationModule
 import LoopVectorization: @turbo, indices
 import ..EquationModule: Node, string_tree
 import ..OperatorEnumModule: OperatorEnum, GenericOperatorEnum
-import ..UtilsModule: @return_on_false, @maybe_turbo, is_bad_array, vals
+import ..UtilsModule: @return_on_false, @maybe_turbo, is_bad_array
 import ..EquationUtilsModule: is_constant
 
 macro return_on_check(val, T, n)
@@ -28,7 +28,7 @@ macro return_on_nonfinite_array(array, T, n)
 end
 
 """
-    eval_tree_array(tree::Node, cX::AbstractMatrix{T}, operators::OperatorEnum; turbo::Bool)
+    eval_tree_array(tree::Node, cX::AbstractMatrix{T}, operators::OperatorEnum; turbo::Bool=false)
 
 Evaluate a binary tree (equation) over a given input data matrix. The
 operators contain all of the operators used. This function fuses doublets
@@ -88,6 +88,7 @@ end
 function _eval_tree_array(
     tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, ::Val{turbo}
 )::Tuple{AbstractVector{T},Bool} where {T<:Real,turbo}
+    n = size(cX, 2)
     # First, we see if there are only constants in the tree - meaning
     # we can just return the constant result.
     if tree.degree == 0
@@ -98,104 +99,88 @@ function _eval_tree_array(
         !flag && return Array{T,1}(undef, size(cX, 2)), false
         return fill(result, size(cX, 2)), true
     elseif tree.degree == 1
+        op = operators.unaops[tree.op]
         if tree.l.degree == 2 && tree.l.l.degree == 0 && tree.l.r.degree == 0
             # op(op2(x, y)), where x, y, z are constants or variables.
-            return deg1_l2_ll0_lr0_eval(
-                tree, cX, vals[tree.op], vals[tree.l.op], operators, Val(turbo)
-            )
+            op_l = operators.binops[tree.l.op]
+            return deg1_l2_ll0_lr0_eval(tree, cX, op, op_l, Val(turbo))
         elseif tree.l.degree == 1 && tree.l.l.degree == 0
             # op(op2(x)), where x is a constant or variable.
-            return deg1_l1_ll0_eval(
-                tree, cX, vals[tree.op], vals[tree.l.op], operators, Val(turbo)
-            )
-        else
-            # op(x), for any x.
-            return deg1_eval(tree, cX, vals[tree.op], operators, Val(turbo))
+            op_l = operators.unaops[tree.l.op]
+            return deg1_l1_ll0_eval(tree, cX, op, op_l, Val(turbo))
         end
+
+        # op(x), for any x.
+        (cumulator, complete) = _eval_tree_array(tree.l, cX, operators, Val(turbo))
+        @return_on_false complete cumulator
+        @return_on_nonfinite_array cumulator T n
+        return deg1_eval(cumulator, op, Val(turbo))
+
     elseif tree.degree == 2
+        op = operators.binops[tree.op]
         # TODO - add op(op2(x, y), z) and op(x, op2(y, z))
+        # op(x, y), where x, y are constants or variables.
         if tree.l.degree == 0 && tree.r.degree == 0
-            # op(x, y), where x, y are constants or variables.
-            return deg2_l0_r0_eval(tree, cX, vals[tree.op], operators, Val(turbo))
-        elseif tree.l.degree == 0
-            # op(x, y), where x is a constant or variable but y is not.
-            return deg2_l0_eval(tree, cX, vals[tree.op], operators, Val(turbo))
+            return deg2_l0_r0_eval(tree, cX, op, Val(turbo))
         elseif tree.r.degree == 0
+            (cumulator_l, complete) = _eval_tree_array(tree.l, cX, operators, Val(turbo))
+            @return_on_false complete cumulator_l
+            @return_on_nonfinite_array cumulator_l T n
             # op(x, y), where y is a constant or variable but x is not.
-            return deg2_r0_eval(tree, cX, vals[tree.op], operators, Val(turbo))
-        else
-            # op(x, y), for any x or y
-            return deg2_eval(tree, cX, vals[tree.op], operators, Val(turbo))
+            return deg2_r0_eval(tree, cumulator_l, cX, op, Val(turbo))
+        elseif tree.l.degree == 0
+            (cumulator_r, complete) = _eval_tree_array(tree.r, cX, operators, Val(turbo))
+            @return_on_false complete cumulator_r
+            @return_on_nonfinite_array cumulator_r T n
+            # op(x, y), where x is a constant or variable but y is not.
+            return deg2_l0_eval(tree, cumulator_r, cX, op, Val(turbo))
         end
+        (cumulator_l, complete) = _eval_tree_array(tree.l, cX, operators, Val(turbo))
+        @return_on_false complete cumulator_l
+        @return_on_nonfinite_array cumulator_l T n
+        (cumulator_r, complete) = _eval_tree_array(tree.r, cX, operators, Val(turbo))
+        @return_on_false complete cumulator_r
+        @return_on_nonfinite_array cumulator_r T n
+        # op(x, y), for any x or y
+        return deg2_eval(cumulator_l, cumulator_r, op, Val(turbo))
     end
 end
 
 function deg2_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,turbo}
-    n = size(cX, 2)
-    (cumulator, complete) = _eval_tree_array(tree.l, cX, operators, Val(turbo))
-    @return_on_false complete cumulator
-    @return_on_nonfinite_array cumulator T n
-    (array2, complete2) = _eval_tree_array(tree.r, cX, operators, Val(turbo))
-    @return_on_false complete2 cumulator
-    @return_on_nonfinite_array array2 T n
-    op = operators.binops[op_idx]
-
-    # We check inputs (and intermediates), not outputs.
-    @maybe_turbo turbo for j in indices(cumulator)
-        x = op(cumulator[j], array2[j])::T
-        cumulator[j] = x
+    cumulator_l::AbstractVector{T}, cumulator_r::AbstractVector{T}, op::F, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
+    @maybe_turbo turbo for j in indices(cumulator_l)
+        x = op(cumulator_l[j], cumulator_r[j])::T
+        cumulator_l[j] = x
     end
-    # return (cumulator, finished_loop) #
-    return (cumulator, true)
+    return (cumulator_l, true)
 end
 
 function deg1_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,turbo}
-    n = size(cX, 2)
-    (cumulator, complete) = _eval_tree_array(tree.l, cX, operators, Val(turbo))
-    @return_on_false complete cumulator
-    @return_on_nonfinite_array cumulator T n
-    op = operators.unaops[op_idx]
+    cumulator::AbstractVector{T}, op::F, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
     @maybe_turbo turbo for j in indices(cumulator)
         x = op(cumulator[j])::T
         cumulator[j] = x
     end
-    return (cumulator, true) #
+    return (cumulator, true)
 end
 
 function deg0_eval(
     tree::Node{T}, cX::AbstractMatrix{T}
 )::Tuple{AbstractVector{T},Bool} where {T<:Real}
-    n = size(cX, 2)
     if tree.constant
+        n = size(cX, 2)
         return (fill(tree.val::T, n), true)
     else
         return (cX[tree.feature, :], true)
     end
 end
 
 function deg1_l2_ll0_lr0_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    ::Val{op_l_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,op_l_idx,turbo}
+    tree::Node{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,F2,turbo}
     n = size(cX, 2)
-    op = operators.unaops[op_idx]
-    op_l = operators.binops[op_l_idx]
     if tree.l.l.constant && tree.l.r.constant
         val_ll = tree.l.l.val::T
         val_lr = tree.l.r.val::T
@@ -243,16 +228,9 @@ end
 
 # op(op2(x)) for x variable or constant
 function deg1_l1_ll0_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    ::Val{op_l_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,op_l_idx,turbo}
+    tree::Node{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,F2,turbo}
     n = size(cX, 2)
-    op = operators.unaops[op_idx]
-    op_l = operators.unaops[op_l_idx]
     if tree.l.l.constant
         val_ll = tree.l.l.val::T
         @return_on_check val_ll T n
@@ -275,14 +253,9 @@ end
 
 # op(x, y) for x and y variable/constant
 function deg2_l0_r0_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,turbo}
+    tree::Node{T}, cX::AbstractMatrix{T}, op::F, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
     n = size(cX, 2)
-    op = operators.binops[op_idx]
     if tree.l.constant && tree.r.constant
         val_l = tree.l.val::T
         @return_on_check val_l T n
@@ -323,17 +296,9 @@ end
 
 # op(x, y) for x variable/constant, y arbitrary
 function deg2_l0_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,turbo}
+    tree::Node{T}, cumulator::AbstractVector{T}, cX::AbstractArray{T}, op::F, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
     n = size(cX, 2)
-    (cumulator, complete) = _eval_tree_array(tree.r, cX, operators, Val(turbo))
-    @return_on_false complete cumulator
-    @return_on_nonfinite_array cumulator T n
-    op = operators.binops[op_idx]
     if tree.l.constant
         val = tree.l.val::T
         @return_on_check val T n
@@ -353,17 +318,9 @@ end
 
 # op(x, y) for x arbitrary, y variable/constant
 function deg2_r0_eval(
-    tree::Node{T},
-    cX::AbstractMatrix{T},
-    ::Val{op_idx},
-    operators::OperatorEnum,
-    ::Val{turbo},
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,turbo}
+    tree::Node{T}, cumulator::AbstractVector{T}, cX::AbstractArray{T}, op::F, ::Val{turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
     n = size(cX, 2)
-    (cumulator, complete) = _eval_tree_array(tree.l, cX, operators, Val(turbo))
-    @return_on_false complete cumulator
-    @return_on_nonfinite_array cumulator T n
-    op = operators.binops[op_idx]
     if tree.r.constant
         val = tree.r.val::T
         @return_on_check val T n
@@ -394,9 +351,9 @@ function _eval_constant_tree(
     if tree.degree == 0
         return deg0_eval_constant(tree)
     elseif tree.degree == 1
-        return deg1_eval_constant(tree, vals[tree.op], operators)
+        return deg1_eval_constant(tree, operators.unaops[tree.op], operators)
     else
-        return deg2_eval_constant(tree, vals[tree.op], operators)
+        return deg2_eval_constant(tree, operators.binops[tree.op], operators)
     end
 end
 
@@ -405,19 +362,17 @@ end
 end
 
 function deg1_eval_constant(
-    tree::Node{T}, ::Val{op_idx}, operators::OperatorEnum
-)::Tuple{T,Bool} where {T<:Real,op_idx}
-    op = operators.unaops[op_idx]
+    tree::Node{T}, op::F, operators::OperatorEnum
+)::Tuple{T,Bool} where {T<:Real,F}
     (cumulator, complete) = _eval_constant_tree(tree.l, operators)
     !complete && return zero(T), false
     output = op(cumulator)::T
     return output, isfinite(output)
 end
 
 function deg2_eval_constant(
-    tree::Node{T}, ::Val{op_idx}, operators::OperatorEnum
-)::Tuple{T,Bool} where {T<:Real,op_idx}
-    op = operators.binops[op_idx]
+    tree::Node{T}, op::F, operators::OperatorEnum
+)::Tuple{T,Bool} where {T<:Real,F}
     (cumulator, complete) = _eval_constant_tree(tree.l, operators)
     !complete && return zero(T), false
     (cumulator2, complete2) = _eval_constant_tree(tree.r, operators)
@@ -442,31 +397,29 @@ function differentiable_eval_tree_array(
             return (cX[tree.feature, :], true)
         end
     elseif tree.degree == 1
-        return deg1_diff_eval(tree, cX, vals[tree.op], operators)
+        return deg1_diff_eval(tree, cX, operators.unaops[tree.op], operators)
     else
-        return deg2_diff_eval(tree, cX, vals[tree.op], operators)
+        return deg2_diff_eval(tree, cX, operators.binops[tree.op], operators)
     end
 end
 
 function deg1_diff_eval(
-    tree::Node{T1}, cX::AbstractMatrix{T}, ::Val{op_idx}, operators::OperatorEnum
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,T1}
+    tree::Node{T1}, cX::AbstractMatrix{T}, op::F, operators::OperatorEnum
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,T1}
     (left, complete) = differentiable_eval_tree_array(tree.l, cX, operators)
     @return_on_false complete left
-    op = operators.unaops[op_idx]
     out = op.(left)
     no_nans = !any(x -> (!isfinite(x)), out)
     return (out, no_nans)
 end
 
 function deg2_diff_eval(
-    tree::Node{T1}, cX::AbstractMatrix{T}, ::Val{op_idx}, operators::OperatorEnum
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,op_idx,T1}
+    tree::Node{T1}, cX::AbstractMatrix{T}, op::F, operators::OperatorEnum
+)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,T1}
     (left, complete) = differentiable_eval_tree_array(tree.l, cX, operators)
     @return_on_false complete left
     (right, complete2) = differentiable_eval_tree_array(tree.r, cX, operators)
     @return_on_false complete2 left
-    op = operators.binops[op_idx]
     out = op.(left, right)
     no_nans = !any(x -> (!isfinite(x)), out)
     return (out, no_nans)
@@ -557,30 +510,32 @@ function _eval_tree_array_generic(
             end
         end
     elseif tree.degree == 1
-        return deg1_eval_generic(tree, cX, vals[tree.op], operators, Val(throw_errors))
+        return deg1_eval_generic(
+            tree, cX, operators.unaops[tree.op], operators, Val(throw_errors)
+        )
     else
-        return deg2_eval_generic(tree, cX, vals[tree.op], operators, Val(throw_errors))
+        return deg2_eval_generic(
+            tree, cX, operators.binops[tree.op], operators, Val(throw_errors)
+        )
     end
 end
 
 function deg1_eval_generic(
-    tree, cX, ::Val{op_idx}, operators::GenericOperatorEnum, ::Val{throw_errors}
-) where {op_idx,throw_errors}
+    tree, cX, op::F, operators::GenericOperatorEnum, ::Val{throw_errors}
+) where {F,throw_errors}
     left, complete = eval_tree_array(tree.l, cX, operators)
     !throw_errors && !complete && return nothing, false
-    op = operators.unaops[op_idx]
     !throw_errors && !hasmethod(op, Tuple{typeof(left)}) && return nothing, false
     return op(left), true
 end
 
 function deg2_eval_generic(
-    tree, cX, ::Val{op_idx}, operators::GenericOperatorEnum, ::Val{throw_errors}
-) where {op_idx,throw_errors}
+    tree, cX, op::F, operators::GenericOperatorEnum, ::Val{throw_errors}
+) where {F,throw_errors}
     left, complete = eval_tree_array(tree.l, cX, operators)
     !throw_errors && !complete && return nothing, false
     right, complete = eval_tree_array(tree.r, cX, operators)
     !throw_errors && !complete && return nothing, false
-    op = operators.binops[op_idx]
     !throw_errors &&
         !hasmethod(op, Tuple{typeof(left),typeof(right)}) &&
         return nothing, false