Add @turbo mode to gradient calculation

MilesCranmer · MilesCranmer · commit 9990e2bbe167 · 2022-11-11T21:05:25.000-05:00
diff --git a/src/EvaluateEquation.jl b/src/EvaluateEquation.jl
@@ -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
diff --git a/src/EvaluateEquationDerivative.jl b/src/EvaluateEquationDerivative.jl
@@ -1,9 +1,9 @@
 module EvaluateEquationDerivativeModule
 
-import LoopVectorization: indices
+import LoopVectorization: indices, @turbo
 import ..EquationModule: Node
 import ..OperatorEnumModule: OperatorEnum
-import ..UtilsModule: @return_on_false2, is_bad_array
+import ..UtilsModule: @return_on_false2, @maybe_turbo, is_bad_array
 import ..EquationUtilsModule: count_constants, index_constants, NodeIndex
 import ..EvaluateEquationModule: deg0_eval
 
@@ -16,7 +16,7 @@ function assert_autodiff_enabled(operators::OperatorEnum)
 end
 
 """
-    eval_diff_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, direction::Int)
+    eval_diff_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, direction::Int; turbo::Bool=false)
 
 Compute the forward derivative of an expression, using a similar
 structure and optimization to eval_tree_array. `direction` is the index of a particular
@@ -30,33 +30,48 @@ respect to `x1`.
 - `operators::OperatorEnum`: The operators used to create the `tree`. Note that `operators.enable_autodiff`
     must be `true`. This is needed to create the derivative operations.
 - `direction::Int`: The index of the variable to take the derivative with respect to.
+- `turbo::Bool`: Use `LoopVectorization.@turbo` for faster evaluation.
 
 # Returns
 
 - `(evaluation, derivative, complete)::Tuple{AbstractVector{T}, AbstractVector{T}, Bool}`: the normal evaluation,
     the derivative, and whether the evaluation completed as normal (or encountered a nan or inf).
 """
 function eval_diff_tree_array(
-    tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, direction::Int
+    tree::Node{T},
+    cX::AbstractMatrix{T},
+    operators::OperatorEnum,
+    direction::Int;
+    turbo::Bool=false,
 )::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real}
     assert_autodiff_enabled(operators)
     # TODO: Implement quick check for whether the variable is actually used
     # in this tree. Otherwise, return zero.
-    return _eval_diff_tree_array(tree, cX, operators, direction)
+    return _eval_diff_tree_array(
+        tree, cX, operators, direction, (turbo ? Val(true) : Val(false))
+    )
 end
 function eval_diff_tree_array(
-    tree::Node{T1}, cX::AbstractMatrix{T2}, operators::OperatorEnum, direction::Int;
+    tree::Node{T1},
+    cX::AbstractMatrix{T2},
+    operators::OperatorEnum,
+    direction::Int;
+    turbo::Bool=false,
 ) where {T1<:Real,T2<:Real}
     T = promote_type(T1, T2)
     @warn "Warning: eval_diff_tree_array received mixed types: tree=$(T1) and data=$(T2)."
     tree = convert(Node{T}, tree)
     cX = convert(AbstractMatrix{T}, cX)
-    return eval_diff_tree_array(tree, cX, operators, direction)
+    return eval_diff_tree_array(tree, cX, operators, direction; turbo=turbo)
 end
 
 function _eval_diff_tree_array(
-    tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, direction::Int
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real}
+    tree::Node{T},
+    cX::AbstractMatrix{T},
+    operators::OperatorEnum,
+    direction::Int,
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,turbo}
     evaluation, derivative, complete = if tree.degree == 0
         diff_deg0_eval(tree, cX, direction)
     elseif tree.degree == 1
@@ -67,6 +82,7 @@ function _eval_diff_tree_array(
             operators.diff_unaops[tree.op],
             operators,
             direction,
+            Val(turbo),
         )
     else
         diff_deg2_eval(
@@ -76,6 +92,7 @@ function _eval_diff_tree_array(
             operators.diff_binops[tree.op],
             operators,
             direction,
+            Val(turbo),
         )
     end
     @return_on_false2 complete evaluation derivative
@@ -99,15 +116,16 @@ function diff_deg1_eval(
     diff_op::dF,
     operators::OperatorEnum,
     direction::Int,
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,F,dF}
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,F,dF,turbo}
     n = size(cX, 2)
     (cumulator, dcumulator, complete) = _eval_diff_tree_array(
-        tree.l, cX, operators, direction
+        tree.l, cX, operators, direction, Val(turbo)
     )
     @return_on_false2 complete cumulator dcumulator
 
     # TODO - add type assertions to get better speed:
-    @inbounds @simd for j in indices((cumulator, dcumulator))
+    @maybe_turbo turbo for j in indices((cumulator, dcumulator))
         x = op(cumulator[j])::T
         dx = diff_op(cumulator[j])::T * dcumulator[j]
 
@@ -124,21 +142,21 @@ function diff_deg2_eval(
     diff_op::dF,
     operators::OperatorEnum,
     direction::Int,
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,F,dF}
-    n = size(cX, 2)
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,F,dF,turbo}
     (cumulator, dcumulator, complete) = _eval_diff_tree_array(
-        tree.l, cX, operators, direction
+        tree.l, cX, operators, direction, Val(turbo)
     )
     @return_on_false2 complete cumulator dcumulator
     (array2, dcumulator2, complete2) = _eval_diff_tree_array(
-        tree.r, cX, operators, direction
+        tree.r, cX, operators, direction, Val(turbo)
     )
     @return_on_false2 complete2 array2 dcumulator2
 
-    @inbounds @simd for j in indices((cumulator, dcumulator, array2, dcumulator2))
-        x = op(cumulator[j], array2[j])
+    @maybe_turbo turbo for j in indices((cumulator, dcumulator, array2, dcumulator2))
+        x = op(cumulator[j], array2[j])::T
 
-        first, second = diff_op(cumulator[j], array2[j])
+        first, second = diff_op(cumulator[j], array2[j])::Tuple{T,T}
         dx = first * dcumulator[j] + second * dcumulator2[j]
 
         cumulator[j] = x
@@ -148,7 +166,7 @@ function diff_deg2_eval(
 end
 
 """
-    eval_grad_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum; variable::Bool=false)
+    eval_grad_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum; variable::Bool=false, turbo::Bool=false)
 
 Compute the forward-mode derivative of an expression, using a similar
 structure and optimization to eval_tree_array. `variable` specifies whether
@@ -163,14 +181,19 @@ to every constant in the expression.
     must be `true`. This is needed to create the derivative operations.
 - `variable::Bool`: Whether to take derivatives with respect to features (i.e., `cX` - with `variable=true`),
     or with respect to every constant in the expression (`variable=false`).
+- `turbo::Bool`: Use `LoopVectorization.@turbo` for faster evaluation.
 
 # Returns
 
 - `(evaluation, gradient, complete)::Tuple{AbstractVector{T}, AbstractMatrix{T}, Bool}`: the normal evaluation,
     the gradient, and whether the evaluation completed as normal (or encountered a nan or inf).
 """
 function eval_grad_tree_array(
-    tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum; variable::Bool=false
+    tree::Node{T},
+    cX::AbstractMatrix{T},
+    operators::OperatorEnum;
+    variable::Bool=false,
+    turbo::Bool=false,
 )::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real}
     assert_autodiff_enabled(operators)
     n = size(cX, 2)
@@ -181,7 +204,14 @@ function eval_grad_tree_array(
     end
     index_tree = index_constants(tree, 0)
     return eval_grad_tree_array(
-        tree, n, n_gradients, index_tree, cX, operators, (variable ? Val(true) : Val(false))
+        tree,
+        n,
+        n_gradients,
+        index_tree,
+        cX,
+        operators,
+        (variable ? Val(true) : Val(false)),
+        (turbo ? Val(true) : Val(false)),
     )
 end
 
@@ -193,20 +223,29 @@ function eval_grad_tree_array(
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
     ::Val{variable},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable}
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable,turbo}
     evaluation, gradient, complete = _eval_grad_tree_array(
-        tree, n, n_gradients, index_tree, cX, operators, Val(variable)
+        tree, n, n_gradients, index_tree, cX, operators, Val(variable), Val(turbo)
     )
     @return_on_false2 complete evaluation gradient
     return evaluation, gradient, !(is_bad_array(evaluation) || is_bad_array(gradient))
 end
 
 function eval_grad_tree_array(
-    tree::Node{T1}, cX::AbstractMatrix{T2}, operators::OperatorEnum; variable::Bool=false
+    tree::Node{T1},
+    cX::AbstractMatrix{T2},
+    operators::OperatorEnum;
+    variable::Bool=false,
+    turbo::Bool=false,
 ) where {T1<:Real,T2<:Real}
     T = promote_type(T1, T2)
     return eval_grad_tree_array(
-        convert(Node{T}, tree), convert(AbstractMatrix{T}, cX), operators; variable=variable
+        convert(Node{T}, tree),
+        convert(AbstractMatrix{T}, cX),
+        operators;
+        variable=variable,
+        turbo=turbo,
     )
 end
 
@@ -218,7 +257,8 @@ function _eval_grad_tree_array(
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
     ::Val{variable},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable}
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable,turbo}
     if tree.degree == 0
         grad_deg0_eval(tree, n, n_gradients, index_tree, cX, Val(variable))
     elseif tree.degree == 1
@@ -232,6 +272,7 @@ function _eval_grad_tree_array(
             operators.diff_unaops[tree.op],
             operators,
             Val(variable),
+            Val(turbo),
         )
     else
         grad_deg2_eval(
@@ -244,6 +285,7 @@ function _eval_grad_tree_array(
             operators.diff_binops[tree.op],
             operators,
             Val(variable),
+            Val(turbo),
         )
     end
 end
@@ -278,13 +320,14 @@ function grad_deg1_eval(
     diff_op::dF,
     operators::OperatorEnum,
     ::Val{variable},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,F,dF,variable}
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,F,dF,variable,turbo}
     (cumulator, dcumulator, complete) = eval_grad_tree_array(
-        tree.l, n, n_gradients, index_tree.l, cX, operators, Val(variable)
+        tree.l, n, n_gradients, index_tree.l, cX, operators, Val(variable), Val(turbo)
     )
     @return_on_false2 complete cumulator dcumulator
 
-    @inbounds @simd for j in indices((cumulator, dcumulator), (1, 2))
+    @maybe_turbo turbo for j in indices((cumulator, dcumulator), (1, 2))
         x = op(cumulator[j])::T
         dx = diff_op(cumulator[j])::T
 
@@ -306,17 +349,18 @@ function grad_deg2_eval(
     diff_op::dF,
     operators::OperatorEnum,
     ::Val{variable},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,F,dF,variable}
+    ::Val{turbo},
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,F,dF,variable,turbo}
     (cumulator1, dcumulator1, complete) = eval_grad_tree_array(
-        tree.l, n, n_gradients, index_tree.l, cX, operators, Val(variable)
+        tree.l, n, n_gradients, index_tree.l, cX, operators, Val(variable), Val(turbo)
     )
     @return_on_false2 complete cumulator1 dcumulator1
     (cumulator2, dcumulator2, complete2) = eval_grad_tree_array(
-        tree.r, n, n_gradients, index_tree.r, cX, operators, Val(variable)
+        tree.r, n, n_gradients, index_tree.r, cX, operators, Val(variable), Val(turbo)
     )
     @return_on_false2 complete2 cumulator1 dcumulator1
 
-    @inbounds @simd for j in indices(
+    @maybe_turbo turbo for j in indices(
         (cumulator1, cumulator2, dcumulator1, dcumulator2), (1, 1, 2, 2)
     )
         c1 = cumulator1[j]
diff --git a/test/test_derivatives.jl b/test/test_derivatives.jl
@@ -6,7 +6,8 @@ using Zygote
 using LinearAlgebra
 
 seed = 0
-pow_abs2(x, y) = abs(x)^y
+# SIMD doesn't like abs(x) ^ y for some reason.
+pow_abs2(x, y) = exp(y * log(abs(x)))
 custom_cos(x) = cos(x)^2
 
 equation1(x1, x2, x3) = x1 + x2 + x3 + 3.2
@@ -35,8 +36,12 @@ function array_test(ar1, ar2; rtol=0.1)
     return isapprox(ar1, ar2; rtol=rtol)
 end
 
-for type in [Float16, Float32, Float64]
-    println("Testing derivatives with respect to variables, with type=$(type).")
+for type in [Float16, Float32, Float64], turbo in [true, false]
+    type == Float16 && turbo && continue
+
+    println(
+        "Testing derivatives with respect to variables, with type=$(type) and turbo=$(turbo).",
+    )
     rng = MersenneTwister(seed)
     nfeatures = 3
     N = 100
@@ -72,10 +77,16 @@ for type in [Float16, Float32, Float64]
         )
         # Convert tuple of vectors to matrix:
         true_grad = reduce(hcat, true_grad)'
-        predicted_grad = eval_grad_tree_array(tree, X, operators; variable=true)[2]
+        predicted_grad = eval_grad_tree_array(
+            tree, X, operators; variable=true, turbo=turbo
+        )[2]
         predicted_grad2 =
             reduce(
-                hcat, [eval_diff_tree_array(tree, X, operators, i)[2] for i in 1:nfeatures]
+                hcat,
+                [
+                    eval_diff_tree_array(tree, X, operators, i; turbo=turbo)[2] for
+                    i in 1:nfeatures
+                ],
             )'
         predicted_grad3 = tree'(X)
 
@@ -98,7 +109,7 @@ for type in [Float16, Float32, Float64]
     local tree
     tree = equation4(nx1, nx2, nx3)
     tree = convert(Node{type}, tree)
-    predicted_grad = eval_grad_tree_array(tree, X, operators; variable=false)[2]
+    predicted_grad = eval_grad_tree_array(tree, X, operators; variable=false, turbo=turbo)[2]
     @test array_test(predicted_grad[1, :], X[1, :])
 
     # More complex expression:
@@ -123,7 +134,7 @@ for type in [Float16, Float32, Float64]
         [X[i, :] for i in 1:nfeatures]...,
     )[1:2]
     true_grad = reduce(hcat, true_grad)'
-    predicted_grad = eval_grad_tree_array(tree, X, operators; variable=false)[2]
+    predicted_grad = eval_grad_tree_array(tree, X, operators; variable=false, turbo=turbo)[2]
 
     @test array_test(predicted_grad, true_grad)
     println("Done.")
