Multithreaded rules application (#66)

* initial attempt at multi-threaded rules

* don't multithread by default

* change cutoff

* account for NaN in fuzz tests

* increment version because this would drop support for pre 1.3

* add some tests

* introduce function barrier in hopes that optimizer wont give up

* fix

* updated rule to use new interface

* typo

* change to thread_subtree_cutoff

* add threaded benchmarks; fixup tests

* additional function barrier

* don't splat

* typos

* fix

* reduce splatting

* shuffle around logic

* disable applyall on recurse

* more disabling of applyall in recurse

Author: MasonProtter

.travis.yml

@@ -8,7 +8,7 @@ branches:
 os:
   - linux
 julia:
-  - 1
+  - 1.3.1
   - 1.4.1
   - nightly
 matrix:

Project.toml

@@ -14,7 +14,7 @@ Combinatorics = "1.0"
 NaNMath = "0.3"
 SpecialFunctions = "0.10"
 TimerOutputs = "0.5"
-julia = "1"
+julia = "1.3"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"

benchmark/benchmarks.jl

@@ -5,7 +5,8 @@ using Random
 
 SUITE = BenchmarkGroup()
 
-@syms a b c
+@syms a b c d; Random.seed!(123);
+
 let r = @rule(~x => ~x), rs = RuleSet([r]),
     acr = @rule(~x::isnumber + ~y => ~y)
 
@@ -37,4 +38,19 @@ let r = @rule(~x => ~x), rs = RuleSet([r]),
     overhead["simplify_no_fixp"]["noop:Int"]  = @benchmarkable simplify(1, fixpoint=false)
     overhead["simplify_no_fixp"]["noop:Sym"]  = @benchmarkable simplify($a, fixpoint=false)
     overhead["simplify_no_fixp"]["noop:Term"] = @benchmarkable simplify($(a+2), fixpoint=false)
+
+    function random_term(len; atoms, funs, fallback_atom=1)
+        xs = rand(atoms, len)
+        while length(xs) > 1
+            xs = map(Iterators.partition(xs, 2)) do xy
+                x = xy[1]; y = get(xy, 2, fallback_atom)
+                rand(funs)(x, y)
+            end
+        end
+        xs[]
+    end
+    ex = random_term(1000, atoms=[a, b, c, d, a^(-1), b^(-1), 1, 2.0], funs=[+, *])
+
+    overhead["simplify_no_fixp"]["randterm:serial"] = @benchmarkable simplify($ex, threaded=false, fixpoint=false)
+    overhead["simplify_no_fixp"]["randterm:thread"] = @benchmarkable simplify($ex, threaded=true,  fixpoint=false)    
 end

src/SymbolicUtils.jl

@@ -27,7 +27,6 @@ else
     end
 end
 
-
 export @syms, term, @fun, showraw
 include("types.jl")
 

src/rule_dsl.jl

@@ -248,19 +248,41 @@ struct RuleRewriteError
     expr
 end
 
-function (r::RuleSet)(term, context=EmptyCtx(); depth=typemax(Int), applyall=false, recurse=true)
+node_count(atom, count; cutoff) = count + 1
+node_count(t::Term, count=0; cutoff=100) = sum(node_count(arg, count; cutoff=cutoff) for arg ∈ arguments(t))
+
+function _recurse_apply_ruleset_threaded(r::RuleSet, term, context; depth, thread_subtree_cutoff)
+    _args = map(arguments(term)) do arg
+        if node_count(arg) > thread_subtree_cutoff
+            Threads.@spawn r(arg, context; depth=depth-1, threaded=true,
+                             thread_subtree_cutoff=thread_subtree_cutoff)
+        else
+            r(arg, context; depth=depth-1, threaded=false)
+        end
+    end
+    args = map(t -> t isa Task ? fetch(t) : t, _args)
+    Term{symtype(term)}(operation(term), args)
+end
+
+function (r::RuleSet)(term, context=EmptyCtx();  depth=typemax(Int), applyall::Bool=false, recurse::Bool=true,
+                      threaded::Bool=false, thread_subtree_cutoff::Int=100)
     rules = r.rules
     term = to_symbolic(term)
     # simplify the subexpressions
     if depth == 0
         return term
     end
     if term isa Symbolic
-        if term isa Term && recurse
-            expr = Term{symtype(term)}(operation(term),
-                                       map(t -> r(t, context, depth=depth-1), arguments(term)))
+        expr = if term isa Term && recurse
+            if threaded
+                _recurse_apply_ruleset_threaded(r, term, context; depth=depth,
+                                                thread_subtree_cutoff=thread_subtree_cutoff)
+            else
+                expr = Term{symtype(term)}(operation(term),
+                                           map(t -> r(t, context, depth=depth-1), arguments(term)))
+            end
         else
-            expr = term
+            term
         end
         for i in 1:length(rules)
             expr′ = try
@@ -282,6 +304,7 @@ function (r::RuleSet)(term, context=EmptyCtx(); depth=typemax(Int), applyall=fal
     return expr # no rule applied
 end
 
+
 getdepth(::RuleSet) = typemax(Int)
 
 function fixpoint(f, x, ctx; kwargs...)

src/simplify.jl

@@ -33,11 +33,11 @@ Applies them once if `fixpoint=false`.
 The `applyall` and `recurse` keywords are forwarded to the enclosed
 `RuleSet`, they are mainly used for internal optimization.
 """
-function simplify(x, ctx=EmptyCtx(); rules=default_rules(x, ctx), fixpoint=true, applyall=true, recurse=true, kwargs...)
+function simplify(x, ctx=EmptyCtx(); rules=default_rules(x, ctx), fixpoint=true, applyall=true, kwargs...)
     if fixpoint
-        SymbolicUtils.fixpoint(rules, x, ctx; recurse=recurse, applyall=recurse)
+        SymbolicUtils.fixpoint(rules, x, ctx; applyall=applyall)
     else
-        rules(x, ctx; recurse=recurse, applyall=recurse)
+        rules(x, ctx; applyall=applyall, kwargs...)
     end
 end
 

test/benchmark.jl

@@ -1,5 +1,5 @@
 # A little trick for travis
-using PkgBenchmark
+using PkgBenchmark, SymbolicUtils
 
 pkgpath = dirname(dirname(pathof(SymbolicUtils)))
 # move it out of the repository so that you can check out different branches

test/fuzz.jl

@@ -1,4 +1,3 @@
-
 include("fuzzlib.jl")
 
 using Random: seed!

test/rulesets.jl

@@ -13,8 +13,7 @@ using SymbolicUtils: fixpoint, getdepth
     ex = 2 * (w+w+α+β)
 
     @eqtest rset(ex) == (((2 * w) + (2 * w)) + (2 * α)) + (2 * β)
-    @eqtest rset(ex) == simplify(ex, rset; fixpoint=false, applyall=false) 
-    
+    @eqtest rset(ex) == simplify(ex; rules=rset, fixpoint=false, applyall=false) 
     @eqtest fixpoint(rset, ex, "ctx") == ((2 * (2 * w)) + (2 * α)) + (2 * β)
 end
 
@@ -103,6 +102,16 @@ end
     @test sprint(io->Base.showerror(io, err)) == "Failed to apply rule ~x + ~(y::pred) => ~x on expression a + b"
 end
 
+@testset "Threading" begin
+    @syms a b c d
+    ex = (((0.6666666666666666 / (c / 1)) + ((1 * a) / (c / 1))) +
+          (1.0 / (((1 * d) / (1 + b)) * (1 / b)))) +
+          ((((1 * a) + (1 * a)) / ((2.0 * (d + 1)) / 1.0)) +
+           ((((d * 1) / (1 + c)) * 2.0) / ((1 / d) + (1 / c))))
+    @eqtest simplify(ex) == simplify(ex, threaded=true, thread_subtree_cutoff=3)
+    @test SymbolicUtils.node_count(a + b * c / d) == 4
+end
+
 @testset "timerwrite" begin
     @syms a b c d
     expr1 = foldr((x,y)->rand([*, /])(x,y), rand([a,b,c,d], 100))