Merge pull request #389 from JuliaSymbolics/ale/egraphs

[WIP] EGraphs optimization

Author: shashi

Project.toml

@@ -36,7 +36,7 @@ DocStringExtensions = "0.8"
 DynamicPolynomials = "0.3"
 IfElse = "0.1"
 LabelledArrays = "1.5"
-Metatheory = "1.1"
+Metatheory = "1.2"
 MultivariatePolynomials = "0.3"
 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)

src/types.jl

@@ -67,6 +67,25 @@ Base.isequal(::Symbolic, ::Symbolic) = false
 
 ### End of interface
 
+### Metatheory.jl e-graph rewriting integration 
+
+"""
+    SymtypeAnalysis
+
+This abstract type is used to identify the EGraph analysis 
+that keeps track of symtype through an EGraph. This must 
+be added to every EGraph that is used in SymbolicUtils.
+"""
+abstract type SymtypeAnalysis <: AbstractAnalysis end
+_getsymtype(T::Type{<:Symbolic{X}}) where X = X
+_getsymtype(T::Type{X}) where {X} = X
+EGraphs.make(an::Type{SymtypeAnalysis}, g::EGraph, n::ENodeLiteral) = symtype(n.value)
+EGraphs.make(an::Type{SymtypeAnalysis}, g::EGraph, n::ENodeTerm{T}) where {T} = _getsymtype(T)
+EGraphs.join(an::Type{SymtypeAnalysis}, A, B) = Union{A, B}
+
+# TODO JOIN egraph analysis
+TermInterface.symtype(ec::EClass) = getdata(ec, SymtypeAnalysis, Any)
+
 function to_symbolic(x)
     Base.depwarn("`to_symbolic(x)` is deprecated, define the interface for your " *
                  "symbolic structure using `istree(x)`, `operation(x)`, `arguments(x)` " *
@@ -348,6 +367,24 @@ function term(f, args...; type = nothing)
     Term{T}(f, [args...])
 end
 
+"""
+    unflatten(t::Symbolic{T})
+Binarizes `Term`s with n-ary operations
+"""
+function unflatten(t::Symbolic{T}) where{T}
+    if istree(t)
+        f = operation(t)
+        if f == (+) || f == (*)   # TODO check out for other n-ary --> binary ops 
+            a = arguments(t)
+            return foldl((x,y) -> Term{T}(f, [x, y]), a)
+        end
+    end
+    return t
+end
+
+unflatten(t) = t
+
+
 """
     similarterm(t, f, args, symtype; metadata=nothing)
 
@@ -366,10 +403,17 @@ different type than `t`, because `f` also influences the result.
 """
 TermInterface.similarterm(t::Type{<:Symbolic}, f, args; metadata=nothing, exprhead=:call) = 
     similarterm(t, f, args, _promote_symtype(f, args); metadata=metadata, exprhead=exprhead)
-    
+
+TermInterface.similarterm(t::Type{<:Symbolic}, f::Symbol, args; metadata=nothing, exprhead=:call) =
+    TermInterface.similarterm(t, eval(f), args; metadata=metadata, exprhead=exprhead)
+
 TermInterface.similarterm(t::Type{<:Term}, f, args, symtype; metadata=nothing, exprhead=:call) = 
     Term{_promote_symtype(f, args)}(f, args; metadata=metadata)
 
+TermInterface.similarterm(t::Type{<:Term}, f::Symbol, args, symtype; metadata=nothing, exprhead=:call) = 
+    Term{_promote_symtype(eval(f), args)}(eval(f), args; metadata=metadata)
+
+
 #--------------------
 #--------------------
 ####  Pretty printing
@@ -549,6 +593,11 @@ showraw(t) = showraw(stdout, t)
 sdict(kv...) = Dict{Any, Number}(kv...)
 
 const SN = Symbolic{<:Number}
+# TODO Reviewme this is necessary for Metatheory.jl egraph rewriting
+# integration. Constructors of `Add, Mul, Pow...` from Base (+, *, ^, ...) 
+# Should now accepts EClasses as arguments. 
+const SN_EC = Union{SN, EClass}
+
 """
     Add(T, coeff, dict::Dict)
 
@@ -583,7 +632,6 @@ end
 
 TermInterface.symtype(a::Add{X}) where {X} = X
 
-
 TermInterface.istree(a::Type{Add}) = true
 
 TermInterface.operation(a::Add) = +
@@ -603,6 +651,17 @@ Base.isequal(a::Add, b::Add) = a.coeff == b.coeff && isequal(a.dict, b.dict)
 
 Base.show(io::IO, a::Add) = show_term(io, a)
 
+function toterm(t::Add{T}) where T
+    args = []
+    for (k, coeff) in t.dict
+        push!(args, coeff == 1 ? k : Term{T}(*, [coeff, k]))
+    end
+    Term{T}(+, args)
+end
+
+toterm(t) = t
+
+
 """
     makeadd(sign, coeff::Number, xs...)
 
@@ -641,7 +700,7 @@ add_t(a,b) = promote_symtype(+, symtype(a), symtype(b))
 sub_t(a,b) = promote_symtype(-, symtype(a), symtype(b))
 sub_t(a) = promote_symtype(-, symtype(a))
 
-function +(a::SN, b::SN)
+function +(a::SN_EC, b::SN_EC)
     if a isa Add
         coeff, dict = makeadd(1, 0, b)
         T = promote_symtype(+, symtype(a), symtype(b))
@@ -652,11 +711,11 @@ function +(a::SN, b::SN)
     Add(add_t(a,b), makeadd(1, 0, a, b)...)
 end
 
-+(a::Number, b::SN) = Add(add_t(a,b), makeadd(1, a, b)...)
++(a::Number, b::SN_EC) = Add(add_t(a,b), makeadd(1, a, b)...)
 
-+(a::SN, b::Number) = Add(add_t(a,b), makeadd(1, b, a)...)
++(a::SN_EC, b::Number) = Add(add_t(a,b), makeadd(1, b, a)...)
 
-+(a::SN) = a
++(a::SN_EC) = a
 
 +(a::Add, b::Add) = Add(add_t(a,b),
                         a.coeff + b.coeff,
@@ -668,17 +727,17 @@ end
 
 -(a::Add) = Add(sub_t(a), -a.coeff, mapvalues((_,v) -> -v, a.dict))
 
--(a::SN) = Add(sub_t(a), makeadd(-1, 0, a)...)
+-(a::SN_EC) = Add(sub_t(a), makeadd(-1, 0, a)...)
 
 -(a::Add, b::Add) = Add(sub_t(a,b),
                         a.coeff - b.coeff,
                         _merge(-, a.dict, b.dict, filter=_iszero))
 
--(a::SN, b::SN) = a + (-b)
+-(a::SN_EC, b::SN_EC) = a + (-b)
 
--(a::Number, b::SN) = a + (-b)
+-(a::Number, b::SN_EC) = a + (-b)
 
--(a::SN, b::Number) = a + (-b)
+-(a::SN_EC, b::Number) = a + (-b)
 
 """
     Mul(T, coeff, dict)
@@ -753,6 +812,16 @@ Base.isequal(a::Mul, b::Mul) = a.coeff == b.coeff && isequal(a.dict, b.dict)
 
 Base.show(io::IO, a::Mul) = show_term(io, a)
 
+function toterm(t::Mul{T}) where T
+    args = []
+    push!(args, t.coeff)
+    for (k, deg) in t.dict
+        push!(args, deg == 1 ? k : Term{T}(^, [k, deg]))
+    end
+    Term{T}(*, args)
+end
+
+
 function makemul(coeff, xs...; d=sdict())
     for x in xs
         if x isa Pow && x.exp isa Number
@@ -777,9 +846,9 @@ end
 mul_t(a,b) = promote_symtype(*, symtype(a), symtype(b))
 mul_t(a) = promote_symtype(*, symtype(a))
 
-*(a::SN) = a
+*(a::SN_EC) = a
 
-function *(a::SN, b::SN)
+function *(a::SN_EC, b::SN_EC)
     # Always make sure Div wraps Mul
     if a isa Div && b isa Div
         Div(a.num * b.num, a.den * b.den)
@@ -796,7 +865,7 @@ end
                         a.coeff * b.coeff,
                         _merge(+, a.dict, b.dict, filter=_iszero))
 
-function *(a::Number, b::SN)
+function *(a::Number, b::SN_EC)
     if iszero(a)
         a
     elseif isone(a)
@@ -812,17 +881,17 @@ function *(a::Number, b::SN)
     end
 end
 
-*(a::SN, b::Number) = b * a
+*(a::SN_EC, b::Number) = b * a
 
-\(a::SN, b::Union{Number, SN}) = b / a
+\(a::SN_EC, b::Union{Number, SN_EC}) = b / a
 
-\(a::Number, b::SN) = b / a
+\(a::Number, b::SN_EC) = b / a
 
-/(a::SN, b::Number) = (b isa Integer ? 1//b : inv(b)) * a
+/(a::SN_EC, b::Number) = (b isa Integer ? 1//b : inv(b)) * a
 
-//(a::Union{SN, Number}, b::SN) = a / b
+//(a::Union{SN_EC, Number}, b::SN_EC) = a / b
 
-//(a::SN, b::T) where {T <: Number} = (one(T) // b) * a
+//(a::SN_EC, b::T) where {T <: Number} = (one(T) // b) * a
 
 """
     Div(numerator_factors, denominator_factors, simplified=false)
@@ -901,7 +970,11 @@ end
 
 Base.show(io::IO, d::Div) = show_term(io, d)
 
-/(a::Union{SN,Number}, b::SN) = Div(a,b)
+function toterm(t::Div{T}) where T
+    Term{T}(/, [t.num, t.den])
+end
+
+/(a::Union{SN_EC,Number}, b::SN_EC) = Div(a,b)
 
 """
     Pow(base, exp)
@@ -944,6 +1017,10 @@ Base.isequal(p::Pow, b::Pow) = isequal(p.base, b.base) && isequal(p.exp, b.exp)
 
 Base.show(io::IO, p::Pow) = show_term(io, p)
 
+function toterm(t::Pow{T}) where T
+    Term{T}(^, [t.base, t.exp])
+end
+
 function makepow(a, b)
     base = a
     exp = b
@@ -954,11 +1031,11 @@ function makepow(a, b)
     return (base, exp)
 end
 
-^(a::SN, b) = Pow(a, b)
+^(a::SN_EC, b) = Pow(a, b)
 
-^(a::SN, b::SN) = Pow(a, b)
+^(a::SN_EC, b::SN_EC) = Pow(a, b)
 
-^(a::Number, b::SN) = Pow(a, b)
+^(a::Number, b::SN_EC) = Pow(a, b)
 
 function ^(a::Mul, b::Number)
     coeff = unstable_pow(a.coeff, b)