Optimize code for composition

Author: Azzaare

Project.toml

@@ -10,13 +10,15 @@ Evolutionary = "86b6b26d-c046-49b6-aa0b-5f0f74682bd6"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ThreadPools = "b189fb0b-2eb5-4ed4-bc0c-d34c51242431"
+Unrolled = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 
 [compat]
 ConstraintDomains = "0.2"
 Dictionaries = "0.3"
 Evolutionary = "0.9"
 OrderedCollections = "1"
 ThreadPools = "2"
+Unrolled = "0.1"
 julia = "1.6"
 
 [extras]

src/CompositionalNetworks.jl

@@ -7,6 +7,7 @@ using Evolutionary
 using OrderedCollections
 using Random
 using ThreadPools
+using Unrolled
 
 # Exports utilities
 export hamming

src/icn.jl

@@ -123,7 +123,19 @@ function _compose(icn::ICN)
     end
 
     l = length(funcs[1])
-    composition = (x; param=nothing, dom_size) -> fill(x, l) .|> map(f -> (y -> f(y; param=param)), funcs[1]) |> funcs[2][1] |> funcs[3][1] |> (y -> funcs[4][1](y; param=param, dom_size=dom_size, nvars=length(x)))
+
+    composition = (x; X=zeros(length(x), l), param=nothing, dom_size) -> if l == 1
+        x |> (y -> funcs[1][1](y; param)) |> funcs[3][1] |>
+        (y -> funcs[4][1](y; param, dom_size, nvars=length(x)))
+    else
+        fill!(@view(X[1:length(x), :]), 0.0)
+        tr_in(Tuple(funcs[1]), X, x, param)
+        for i in 1:length(x)
+            X[i,1] = funcs[2][1](@view X[i,:])
+        end
+        funcs[3][1](@view X[:, 1]) |> (y -> funcs[4][1](y; param, dom_size, nvars=length(x)))
+    end
+
     return composition, symbols
 end
 

src/learn.jl

@@ -111,13 +111,24 @@ function compose_to_string(symbols, name)
     ag = reduce_symbols(symbols[3], ", ", false; prefix=CN * "ag_")
     co = reduce_symbols(symbols[4], ", ", false; prefix=CN * "co_")
 
-    julia_string = """
-    function $name(x; param=nothing, dom_size)
-        fill(x, $tr_length) .|> map(f -> (y -> f(y; param=param)), $tr) |> $ar |> $ag |> (y -> $co(y; param=param, dom_size=dom_size, nvars=length(x)))
+    return if tr_length == 1
+        """
+        function $name(x; X = zeros(length(x), $tr_length), param=nothing, dom_size)
+            x |> (y -> $tr[1](y; param)) |> $ag |> (y -> $co(y; param, dom_size, nvars=length(x)))
+        end
+        """
+    else
+        """
+        function $name(x; X = zeros(length(x), $tr_length), param=nothing, dom_size)
+            fill!(@view(X[1:length(x), :]), 0.0)
+            $(CN)tr_in(Tuple($tr), X, x, param)
+            for i in 1:length(x)
+                X[i,1] = $ar(@view X[i,:])
+            end
+            return $ag(@view X[:, 1]) |> (y -> $co(y; param, dom_size, nvars=length(x)))
+        end
+        """
     end
-    """
-
-    return julia_string
 end
 
 """

src/utils.jl

@@ -61,3 +61,17 @@ end
 function incsert!(d::Dictionary, ind)
     set!(d, ind, isassigned(d, ind) ? d[ind] + 1 : 1)
 end
+
+@unroll function tr_in(tr, X, x, param)
+    @unroll for i in 1:length(tr)
+        X[:,i] = tr[i](x; param)
+    end
+end
+
+# TODO: look for a length limit that make it slow or space-comsuming
+# TODO: handle SMatrix
+# @unroll function ar_in(ar, X, x)
+#     @unroll for i in 1:length(x)
+#         X[i, 1] = ar(@view X[i, :])
+#     end
+# end