Merge branch 'master' into compathelper/new_version/2021-10-24-18-48-57-857-02801455533

Author: shashi

Project.toml

@@ -1,7 +1,7 @@
 name = "SymbolicUtils"
 uuid = "d1185830-fcd6-423d-90d6-eec64667417b"
 authors = ["Shashi Gowda"]
-version = "0.17.0"
+version = "0.18.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -36,7 +36,7 @@ DocStringExtensions = "0.8"
 DynamicPolynomials = "0.3"
 IfElse = "0.1"
 LabelledArrays = "1.5"
-Metatheory = "1.0"
+Metatheory = "1.2"
 MultivariatePolynomials = "0.3, 0.4"
 NaNMath = "0.3"
 Setfield = "0.7, 0.8"

src/SymbolicUtils.jl

@@ -57,8 +57,13 @@ include("simplify_rules.jl")
 export simplify, substitute
 include("api.jl")
 
+# EGraph rewriting
+include("egraph.jl")
+export optimize
+
 include("code.jl")
 
+
 # ADjoints
 include("adjoints.jl")
 

src/egraph.jl

@@ -0,0 +1,101 @@
+using Metatheory.Rewriters
+
+function EGraphs.preprocess(t::Symbolic)
+    toterm(unflatten(t)) 
+end
+
+function symbolicegraph(ex)
+    g = EGraph(ex)
+    analyze!(g, SymbolicUtils.SymtypeAnalysis)
+    settermtype!(g, Term{symtype(ex), Any})
+    return g
+end
+
+
+"""
+Equational rewrite rules for optimizing expressions
+"""
+opt_theory = @theory a b x y  begin
+    a + b == b + a
+    a * b == b * a
+    a * x + a * y == a*(x+y)
+    -1 * a == -a
+    a + (-1 * b) == a - b
+    x^-1 == 1/x 
+    1/x * a == a/x
+    # fraction rules 
+    # (a/b) + (c/b) => (a+c)/b
+    # trig functions
+    sin(x)/cos(x) == tan(x)
+    cos(x)/sin(x) == cot(x)
+    sin(x)^2 + cos(x)^2 --> 1
+    sin(2a) == 2sin(a)cos(a)
+end
+
+
+"""
+Approximation of costs of operators in number 
+of CPU cycles required for the numerical computation
+
+See 
+ * https://latkin.org/blog/2014/11/09/a-simple-benchmark-of-various-math-operations/
+ * https://streamhpc.com/blog/2012-07-16/how-expensive-is-an-operation-on-a-cpu/
+ * https://github.com/triscale-innov/GFlops.jl
+"""
+const op_costs = Dict(
+    (+)     => 1,
+    (-)     => 1,
+    abs     => 2,
+    (*)     => 3,
+    exp     => 18,
+    (/)     => 24,
+    (^)     => 100,
+    log1p   => 124,
+    deg2rad => 125,
+    rad2deg => 125,
+    acos    => 127,
+    asind   => 128,
+    acsch   => 133,
+    sin     => 134,
+    cos     => 134,
+    atan    => 135,
+    tan     => 156,
+)
+# TODO some operator costs are in FLOP and not in cycles!!
+
+function costfun(n::ENodeTerm, g::EGraph, an)
+    op = operation(n)
+    cost = 0
+    cost += get(op_costs, op, 1)
+
+    for id ∈ n.args
+        eclass = g[id]
+        !hasdata(eclass, an) && (cost += Inf; break)
+        cost += last(getdata(eclass, an))
+    end
+    cost
+end
+
+costfun(n::ENodeLiteral, g::EGraph, an) = 0
+
+egraph_simterm(x, head, args, symtype=nothing; metadata=nothing, exprhead=exprhead(x)) = 
+TermInterface.similarterm(typeof(x), head, args, symtype; metadata=metadata, exprhead=exprhead)
+
+
+# Custom similarterm to use in EGraphs on <:Symbolic types that treats everything as a Term 
+function egraph_simterm(x::Type{<:Term}, f, args, symtype=nothing; metadata=nothing, exprhead=:call)
+    T = symtype
+    if T === nothing
+        T = _promote_symtype(f, args)
+    end
+    res = Term{T}(f isa Symbol ? eval(f) : f, args; metadata=metadata);
+    return res
+end 
+
+function optimize(ex; params=SaturationParams(timeout=20))
+    # @show ex
+    g = symbolicegraph(ex)
+    params.simterm = egraph_simterm
+    saturate!(g, opt_theory, params)
+    return extract!(g, costfun)
+end

src/polyform.jl

@@ -236,6 +236,10 @@ function unpolyize(x)
     Postwalk(identity, similarterm=simterm)(x)
 end
 
+function toterm(x::PolyForm)
+    toterm(unpolyize(x))
+end
+
 ## Rational Polynomial form with Div
 
 function polyform_factors(d, pvar2sym, sym2term)