Incorporate LoopVectorization.@turbo into evaluation scheme

Author: MilesCranmer

Project.toml

@@ -5,6 +5,7 @@ version = "0.3.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"

src/EvaluateEquation.jl

@@ -1,5 +1,6 @@
 module EvaluateEquationModule
 
+import LoopVectorization: @turbo, indices
 import ..EquationModule: Node, string_tree
 import ..OperatorEnumModule: OperatorEnum, GenericOperatorEnum
 import ..UtilsModule: @return_on_false, is_bad_array, vals
@@ -132,8 +133,8 @@ function deg2_eval(
     op = operators.binops[op_idx]
 
     # We check inputs (and intermediates), not outputs.
-    @inbounds @simd for j in 1:n
-        x = op(cumulator[j], array2[j])::T
+    @turbo for j in indices(cumulator)
+        x = op(cumulator[j], array2[j])
         cumulator[j] = x
     end
     # return (cumulator, finished_loop) #
@@ -148,8 +149,8 @@ function deg1_eval(
     @return_on_false complete cumulator
     @return_on_nonfinite_array cumulator T n
     op = operators.unaops[op_idx]
-    @inbounds @simd for j in 1:n
-        x = op(cumulator[j])::T
+    @turbo for j in indices(cumulator)
+        x = op(cumulator[j])
         cumulator[j] = x
     end
     return (cumulator, true) #
@@ -191,9 +192,9 @@ function deg1_l2_ll0_lr0_eval(
         @return_on_check val_ll T n
         feature_lr = tree.l.r.feature
         cumulator = Array{T,1}(undef, n)
-        @inbounds @simd for j in 1:n
-            x_l = op_l(val_ll, cX[feature_lr, j])::T
-            x = isfinite(x_l) ? op(x_l)::T : T(Inf) # These will get discovered by _eval_tree_array at end.
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x_l = op_l(val_ll, cX[feature_lr, j])
+            x = isfinite(x_l) ? op(x_l) : T(Inf) # These will get discovered by _eval_tree_array at end.
             cumulator[j] = x
         end
         return (cumulator, true)
@@ -202,19 +203,19 @@ function deg1_l2_ll0_lr0_eval(
         val_lr = tree.l.r.val::T
         @return_on_check val_lr T n
         cumulator = Array{T,1}(undef, n)
-        @inbounds @simd for j in 1:n
-            x_l = op_l(cX[feature_ll, j], val_lr)::T
-            x = isfinite(x_l) ? op(x_l)::T : T(Inf)
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x_l = op_l(cX[feature_ll, j], val_lr)
+            x = isfinite(x_l) ? op(x_l) : T(Inf)
             cumulator[j] = x
         end
         return (cumulator, true)
     else
         feature_ll = tree.l.l.feature
         feature_lr = tree.l.r.feature
         cumulator = Array{T,1}(undef, n)
-        @inbounds @simd for j in 1:n
-            x_l = op_l(cX[feature_ll, j], cX[feature_lr, j])::T
-            x = isfinite(x_l) ? op(x_l)::T : T(Inf)
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x_l = op_l(cX[feature_ll, j], cX[feature_lr, j])
+            x = isfinite(x_l) ? op(x_l) : T(Inf)
             cumulator[j] = x
         end
         return (cumulator, true)
@@ -243,9 +244,9 @@ function deg1_l1_ll0_eval(
     else
         feature_ll = tree.l.l.feature
         cumulator = Array{T,1}(undef, n)
-        @inbounds @simd for j in 1:n
-            x_l = op_l(cX[feature_ll, j])::T
-            x = isfinite(x_l) ? op(x_l)::T : T(Inf)
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x_l = op_l(cX[feature_ll, j])
+            x = isfinite(x_l) ? op(x_l) : T(Inf)
             cumulator[j] = x
         end
         return (cumulator, true)
@@ -270,25 +271,25 @@ function deg2_l0_r0_eval(
         val_l = tree.l.val::T
         @return_on_check val_l T n
         feature_r = tree.r.feature
-        @inbounds @simd for j in 1:n
-            x = op(val_l, cX[feature_r, j])::T
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x = op(val_l, cX[feature_r, j])
             cumulator[j] = x
         end
     elseif tree.r.constant
         cumulator = Array{T,1}(undef, n)
         feature_l = tree.l.feature
         val_r = tree.r.val::T
         @return_on_check val_r T n
-        @inbounds @simd for j in 1:n
-            x = op(cX[feature_l, j], val_r)::T
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x = op(cX[feature_l, j], val_r)
             cumulator[j] = x
         end
     else
         cumulator = Array{T,1}(undef, n)
         feature_l = tree.l.feature
         feature_r = tree.r.feature
-        @inbounds @simd for j in 1:n
-            x = op(cX[feature_l, j], cX[feature_r, j])::T
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x = op(cX[feature_l, j], cX[feature_r, j])
             cumulator[j] = x
         end
     end
@@ -306,14 +307,14 @@ function deg2_l0_eval(
     if tree.l.constant
         val = tree.l.val::T
         @return_on_check val T n
-        @inbounds @simd for j in 1:n
-            x = op(val, cumulator[j])::T
+        @turbo for j in indices(cumulator)
+            x = op(val, cumulator[j])
             cumulator[j] = x
         end
     else
         feature = tree.l.feature
-        @inbounds @simd for j in 1:n
-            x = op(cX[feature, j], cumulator[j])::T
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x = op(cX[feature, j], cumulator[j])
             cumulator[j] = x
         end
     end
@@ -331,14 +332,14 @@ function deg2_r0_eval(
     if tree.r.constant
         val = tree.r.val::T
         @return_on_check val T n
-        @inbounds @simd for j in 1:n
-            x = op(cumulator[j], val)::T
+        @turbo for j in indices(cumulator)
+            x = op(cumulator[j], val)
             cumulator[j] = x
         end
     else
         feature = tree.r.feature
-        @inbounds @simd for j in 1:n
-            x = op(cumulator[j], cX[feature, j])::T
+        @turbo for j in indices((cX, cumulator), (2, 1))
+            x = op(cumulator[j], cX[feature, j])
             cumulator[j] = x
         end
     end