Test turbo in evaluation with different types

Author: MilesCranmer

test/test_evaluation.jl

@@ -8,50 +8,42 @@ operators = OperatorEnum(;
     default_params..., binary_operators=(+, *, /, -), unary_operators=(cos, sin)
 )
 
-# Here, we unittest the fast function evaluation scheme
-# We need to trigger all possible fused functions, with all their logic.
-# These are as follows:
-
-## We fuse (and compile) the following:
-##  - op(op2(x, y)), where x, y, z are constants or variables.
-##  - op(op2(x)), where x is a constant or variable.
-##  - op(x), for any x.
-## We fuse (and compile) the following:
-##  - op(x, y), where x, y are constants or variables.
-##  - op(x, y), where x is a constant or variable but y is not.
-##  - op(x, y), where y is a constant or variable but x is not.
-##  - op(x, y), for any x or y
-for fnc in [
+functions = [
     # deg2_l0_r0_eval
     (x1, x2, x3) -> x1 * x2,
-    (x1, x2, x3) -> x1 * 3.0f0,
-    (x1, x2, x3) -> 3.0f0 * x2,
-    (((x1, x2, x3) -> 3.0f0 * 6.0f0), ((x1, x2, x3) -> Node(; val=3.0f0) * 6.0f0)),
+    (x1, x2, x3) -> x1 * 3.0,
+    (x1, x2, x3) -> 3.0 * x2,
+    (((x1, x2, x3) -> 3.0 * 6.0), ((x1, x2, x3) -> Node(; val=3.0) * 6.0)),
     # deg2_l0_eval
     (x1, x2, x3) -> x1 * sin(x2),
-    (x1, x2, x3) -> 3.0f0 * sin(x2),
+    (x1, x2, x3) -> 3.0 * sin(x2),
 
     # deg2_r0_eval
     (x1, x2, x3) -> sin(x1) * x2,
-    (x1, x2, x3) -> sin(x1) * 3.0f0,
+    (x1, x2, x3) -> sin(x1) * 3.0,
 
     # deg1_l2_ll0_lr0_eval
     (x1, x2, x3) -> cos(x1 * x2),
-    (x1, x2, x3) -> cos(x1 * 3.0f0),
-    (x1, x2, x3) -> cos(3.0f0 * x2),
+    (x1, x2, x3) -> cos(x1 * 3.0),
+    (x1, x2, x3) -> cos(3.0 * x2),
     (
-        ((x1, x2, x3) -> cos(3.0f0 * -0.5f0)),
-        ((x1, x2, x3) -> cos(Node(2, Node(; val=3.0f0), Node(; val=-0.5f0)))),
+        ((x1, x2, x3) -> cos(3.0 * -0.5)),
+        ((x1, x2, x3) -> cos(Node(2, Node(; val=3.0), Node(; val=-0.5)))),
     ),
 
     # deg1_l1_ll0_eval
     (x1, x2, x3) -> cos(sin(x1)),
-    (((x1, x2, x3) -> cos(sin(3.0f0))), ((x1, x2, x3) -> cos(sin(Node(; val=3.0f0))))),
+    (((x1, x2, x3) -> cos(sin(3.0))), ((x1, x2, x3) -> cos(sin(Node(; val=3.0))))),
 
     # everything else:
-    (x1, x2, x3) -> (sin(cos(sin(cos(x1) * x3) * 3.0f0) * -0.5f0) + 2.0f0) * 5.0f0,
+    (x1, x2, x3) -> (sin(cos(sin(cos(x1) * x3) * 3.0) * -0.5) + 2.0) * 5.0,
 ]
 
+for turbo in [false, true], T in [Float16, Float32, Float64], fnc in functions
+
+    # Float16 not implemented:
+    turbo && T == Float16 && continue
+
     # check if fnc is tuple
     if typeof(fnc) <: Tuple
         realfnc = fnc[1]
@@ -61,40 +53,49 @@ for fnc in [
         nodefnc = fnc
     end
 
-    global tree = nodefnc(Node("x1"), Node("x2"), Node("x3"))
+    local tree, X
+    tree = nodefnc(Node("x1"), Node("x2"), Node("x3"))
+    tree = convert(Node{T}, tree)
 
     N = 100
     nfeatures = 3
-    X = randn(MersenneTwister(0), Float32, nfeatures, N)
+    X = randn(MersenneTwister(0), T, nfeatures, N)
 
-    test_y = eval_tree_array(tree, X, operators)[1]
+    test_y = eval_tree_array(tree, X, operators; turbo=turbo)[1]
     true_y = realfnc.(X[1, :], X[2, :], X[3, :])
 
-    zero_tolerance = 1e-6
+    zero_tolerance = (T == Float16 ? 1e-4 : 1e-6)
     @test all(abs.(test_y .- true_y) / N .< zero_tolerance)
 end
 
-# Test specific branches of evaluation code:
-# op(op(<constant>))
-tree = Node(1, Node(1, Node(; val=3.0f0)))
-@test repr(tree) == "cos(cos(3.0))"
-truth = cos(cos(3.0f0))
-@test DynamicExpressions.EvaluateEquationModule.deg1_l1_ll0_eval(
-    tree, [0.0f0]', Val(1), Val(1), operators
-)[1][1] ≈ truth
-
-# op(<constant>, <constant>)
-tree = Node(1, Node(; val=3.0f0), Node(; val=4.0f0))
-@test repr(tree) == "(3.0 + 4.0)"
-truth = 3.0f0 + 4.0f0
-@test DynamicExpressions.EvaluateEquationModule.deg2_l0_r0_eval(
-    tree, [0.0f0]', Val(1), operators
-)[1][1] ≈ truth
-
-# op(op(<constant>, <constant>))
-tree = Node(1, Node(1, Node(; val=3.0f0), Node(; val=4.0f0)))
-@test repr(tree) == "cos(3.0 + 4.0)"
-truth = cos(3.0f0 + 4.0f0)
-@test DynamicExpressions.EvaluateEquationModule.deg1_l2_ll0_lr0_eval(
-    tree, [0.0f0]', Val(1), Val(1), operators
-)[1][1] ≈ truth
+for turbo in [false, true], T in [Float16, Float32, Float64]
+    turbo && T == Float16 && continue
+    # Test specific branches of evaluation code:
+    # op(op(<constant>))
+    local tree
+    tree = Node(1, Node(1, Node(; val=3.0f0)))
+    @test repr(tree) == "cos(cos(3.0))"
+    tree = convert(Node{T}, tree)
+    truth = cos(cos(T(3.0f0)))
+    @test DynamicExpressions.EvaluateEquationModule.deg1_l1_ll0_eval(
+        tree, [zero(T)]', Val(1), Val(1), operators, Val(turbo)
+    )[1][1] ≈ truth
+
+    # op(<constant>, <constant>)
+    tree = Node(1, Node(; val=3.0f0), Node(; val=4.0f0))
+    @test repr(tree) == "(3.0 + 4.0)"
+    tree = convert(Node{T}, tree)
+    truth = T(3.0f0) + T(4.0f0)
+    @test DynamicExpressions.EvaluateEquationModule.deg2_l0_r0_eval(
+        tree, [zero(T)]', Val(1), operators, Val(turbo)
+    )[1][1] ≈ truth
+
+    # op(op(<constant>, <constant>))
+    tree = Node(1, Node(1, Node(; val=3.0f0), Node(; val=4.0f0)))
+    @test repr(tree) == "cos(3.0 + 4.0)"
+    tree = convert(Node{T}, tree)
+    truth = cos(T(3.0f0) + T(4.0f0))
+    @test DynamicExpressions.EvaluateEquationModule.deg1_l2_ll0_lr0_eval(
+        tree, [zero(T)]', Val(1), Val(1), operators, Val(turbo)
+    )[1][1] ≈ truth
+end