Move entire operator loop to code generation

julia/sr.jl

@@ -285,9 +285,7 @@ function evalTreeArray(tree::Node, cX::Array{Float32, 2})::Union{Array{Float32,
             return nothing
         end
         op_idx = tree.op
-        @inbounds @simd for i=1:clen
-            cumulator[i] = UNAOP(op_idx, cumulator[i])
-        end
+        UNAOP!(op_idx, cumulator)
         @inbounds for i=1:clen
             if isinf(cumulator[i]) || isnan(cumulator[i])
                 return nothing
@@ -303,12 +301,8 @@ function evalTreeArray(tree::Node, cX::Array{Float32, 2})::Union{Array{Float32,
         if array2 === nothing
             return nothing
         end
-
         op_idx = tree.op
-
-        @inbounds @simd for i=1:clen
-            cumulator[i] = BINOP(op_idx, cumulator[i], array2[i])
-        end
+        BINOP!(op_idx, cumulator, array2)
         @inbounds for i=1:clen
             if isinf(cumulator[i]) || isnan(cumulator[i])
                 return nothing

pysr/sr.py

@@ -286,27 +286,27 @@ const limitPowComplexity = {"true" if limitPowComplexity else "false"}
 
     op_runner = ""
     if len(binary_operators) > 0:
-        op_runner += f"""
-@inline function BINOP(i::Int, x::Float32, y::Float32)::Float32
+        op_runner += """
+function BINOP!(i::Int, x::Array{Float32, 1}, y::Array{Float32, 1})::Array{Float32, 1}
     if i == 1
-        return @fastmath {binary_operators[0]}(x, y)"""
+        x .= @fastmath """f"{binary_operators[0]}"".(x, y)"
         for i in range(1, len(binary_operators)):
             op_runner += f"""
     elseif i == {i+1}
-        return @fastmath {binary_operators[i]}(x, y)"""
+        x .= @fastmath {binary_operators[i]}.(x, y)"""
         op_runner += """
     end
 end"""
 
     if len(unary_operators) > 0:
-        op_runner += f"""
-@inline function UNAOP(i::Int, x::Float32)::Float32
+        op_runner += """
+function UNAOP!(i::Int, x::Array{Float32, 1})::Array{Float32, 1}
     if i == 1
-        return @fastmath {unary_operators[0]}(x)"""
+        x .= @fastmath """f"{unary_operators[0]}.(x)"
         for i in range(1, len(unary_operators)):
-            op_runner += f"""
+            op_runner += """
     elseif i == {i+1}
-        return @fastmath {unary_operators[i]}(x)"""
+        x .= @fastmath """f"{unary_operators[i]}.(x)"
         op_runner += """
     end
 end"""