refactor: improve division performance

Author: AayushSabharwal

src/polyform.jl

@@ -239,8 +239,11 @@ function quick_cancel(d::BasicSymbolic{T})::BasicSymbolic{T} where {T}
 end
 
 function quick_cancel(x::S, y::S)::Tuple{S, S} where {T <: SymVariant, S <: BasicSymbolic{T}}
+    isequal(x, y) && return one_of_vartype(T), one_of_vartype(T)
     opx = iscall(x) ? operation(x) : nothing
     opy = iscall(y) ? operation(y) : nothing
+    icx = isconst(x)
+    icy = isconst(y)
     if opx === (^) && opy === (^)
         return quick_powpow(x, y)
     elseif opx === (*) && opy === (^)
@@ -249,16 +252,14 @@ function quick_cancel(x::S, y::S)::Tuple{S, S} where {T <: SymVariant, S <: Basi
         return reverse(quick_mulpow(y, x))
     elseif opx === (*) && opy === (*)
         return quick_mulmul(x, y)
-    elseif opx === (^) && !isconst(y)
+    elseif opx === (^) && !icy
         return quick_pow(x, y)
-    elseif opy === (^) && !isconst(x)
+    elseif opy === (^) && !icx
         return reverse(quick_pow(y, x))
-    elseif opx === (*) && !isconst(y)
+    elseif opx === (*) && !icy
         return quick_mul(x, y)
-    elseif opy === (*) && !isconst(x)
+    elseif opy === (*) && !icx
         return reverse(quick_mul(y, x))
-    elseif isequal(x, y)
-        return one_of_vartype(T), one_of_vartype(T)
     else
         return x, y
     end
@@ -340,20 +341,31 @@ end
 
 # Double mul case
 function quick_mulmul(x::S, y::S)::Tuple{S, S} where {T <: SymVariant, S <: BasicSymbolic{T}}
-    yargs = arguments(y)
-    for (i, arg) in enumerate(yargs)
-        newx, newarg = quick_cancel(x, arg)
-        isequal(arg, newarg) && continue
-        if yargs isa ROArgsT
-            yargs = copy(parent(yargs))
+    @match (x, y) begin
+        (BSImpl.AddMul(; coeff = c1, dict = d1, type = t1, shape = s1, variant = vr1), BSImpl.AddMul(; coeff = c2, dict = d2, type = t2, shape = s2, variant = vr2)) => begin
+            newd1 = d1
+            newd2 = d2
+            for (k1, v1) in d1
+                haskey(d2, k1) || continue
+                v2 = d2[k1]
+                if newd1 === d1
+                    newd1 = copy(d1)
+                    newd2 = copy(d2)
+                end
+                delete!((v1 >= v2) ? newd2 : newd1, k1)
+                setindex!((v1 >= v2) ? newd1 : newd2, abs(v1 - v2), k1)
+            end
+            if newd1 === d1
+                return x, y
+            end
+            filter!(!iszero ∘ last, newd1)
+            filter!(!iszero ∘ last, newd2)
+            xx = Mul{T}(c1, newd1; type = t1, shape = s1)
+            yy = Mul{T}(c2, newd2; type = t2, shape = s2)
+
+            return xx, yy
         end
-        yargs[i] = Const{T}(newarg)
-        x = newx
-    end
-    if yargs isa ROArgsT
-        return x, y
-    else
-        return x, mul_worker(T, yargs)
+        _ => _unreachable()
     end
 end
 

src/types.jl

@@ -1092,6 +1092,7 @@ end
 end
 
 @inline function BSImpl.AddMul{T}(coeff, dict, variant::AddMulVariant.T; metadata = nothing, type, shape = default_shape(type), unsafe = false) where {T}
+    @nospecialize coeff
     metadata = parse_metadata(metadata)
     shape = parse_shape(shape)
     dict = parse_dict(T, dict)
@@ -1156,6 +1157,7 @@ struct ArrayOp{T} end
 end
 
 @inline function Add{T}(coeff, dict; kw...) where {T}
+    @nospecialize coeff kw
     coeff = unwrap(coeff)
     dict = unwrap_dict(dict)
     if isempty(dict)
@@ -1173,6 +1175,7 @@ end
 end
 
 @inline function Mul{T}(coeff, dict; kw...) where {T}
+    @nospecialize coeff kw
     coeff = unwrap(coeff)
     dict = unwrap_dict(dict)
     if isempty(dict)
@@ -1190,7 +1193,7 @@ end
         k, v = first(dict)
         if _isone(v)
             @match k begin
-                BSImpl.AddMul(; coeff = c2, dict = d2, variant) && if variant == AddMulVariant.ADD end => begin
+                BSImpl.AddMul(; coeff = c2, dict = d2, variant) && if variant === AddMulVariant.ADD end => begin
                     empty!(dict)
                     for (k, v) in d2
                         dict[k] = -v
@@ -1231,31 +1234,31 @@ end
 Simplify the coefficients of `n` and `d` (numerator and denominator).
 """
 function simplify_coefficients(n, d)
-    if safe_isinteger(n)
-        n = Int(n)
+    return n, d
+end
+
+function safe_div(a::Number, b::Number)::Number
+    # @nospecialize a b
+    if (!(a isa Integer) && safe_isinteger(a))
+        a = Int(a)
     end
-    if safe_isinteger(d)
-        d = Int(d)
+    if (!(b isa Integer) && safe_isinteger(b))
+        b = Int(b)
     end
-    nrat, nc = ratcoeff(n)
-    drat, dc = ratcoeff(d)
-    nrat && drat || return n, d
-    g = gcd(nc, dc) * sign(dc) # make denominator positive
-    invdc = isone(g) ? g : (1 // g)
-    n = maybe_integer(invdc * n)
-    d = maybe_integer(invdc * d)
-
-    return n, d
+    if a isa Integer && b isa Integer
+        return a // b
+    end
+    return a / b
 end
 
 """
     $(TYPEDSIGNATURES)
 """
 function Div{T}(n, d, simplified; type = promote_symtype(/, symtype(n), symtype(d)), kw...) where {T}
-    n = unwrap(n)
-    d = unwrap(d)
+    n = Const{T}(unwrap(n))
+    d = Const{T}(unwrap(d))
 
-    if !(type <: Number)
+    if !_numeric_or_arrnumeric_type(type)
         _iszero(n) && return Const{T}(n)
         _isone(d) && return Const{T}(n)
         return BSImpl.Div{T}(n, d, simplified; type, kw...)
@@ -1264,8 +1267,6 @@ function Div{T}(n, d, simplified; type = promote_symtype(/, symtype(n), symtype(
     if !(T === SafeReal)
         n, d = quick_cancel(Const{T}(n), Const{T}(d))
     end
-    n = unwrap_const(n)
-    d = unwrap_const(d)
 
     _iszero(n) && return Const{T}(n)
     _isone(d) && return Const{T}(n)
@@ -1279,10 +1280,42 @@ function Div{T}(n, d, simplified; type = promote_symtype(/, symtype(n), symtype(
         return Div{T}(n * d.den, d.num, simplified; type, kw...)
     end
 
-    d isa Number && _isone(-d) && return Const{T}(-n)
-    n isa Rat && d isa Rat && return Const{T}(n // d)
+    isconst(d) && _isone(-d) && return Const{T}(-n)
+    if isconst(n) && isconst(d)
+        nn = unwrap_const(n)
+        dd = unwrap_const(d)
+        nn isa Rat && dd isa Rat && return Const{T}(nn // dd)
+        return Const{T}(nn / dd)
+    end
 
-    n, d = simplify_coefficients(n, d)
+    @match (n, d) begin
+        (BSImpl.Const(; val = v1), BSImpl.Const(; val = v2)) => Const{T}(safe_div(v1, v2))
+        (BSImpl.Const(; val = c1), BSImpl.AddMul(; coeff = c2, dict, type, shape, variant)) && if variant == AddMulVariant.MUL end => begin
+            if c1 isa Rat && c2 isa Rat
+                tmp = c1 // c2
+                c1 = tmp.num
+                c2 = tmp.den
+            end
+            n = Const{T}(c1)
+            d = Mul{T}(c2, dict, ; type, shape)
+        end
+        (BSImpl.AddMul(; coeff, dict, type, shape, variant), BSImpl.Const(; val)) && if variant == AddMulVariant.MUL end => begin
+            return Mul{T}(safe_div(coeff, val), dict, ; type, shape)
+        end
+        (BSImpl.AddMul(; coeff = c1, dict = d1, type = t1, shape = sh1, variant = v1),
+            BSImpl.AddMul(; coeff = c2, dict = d2, type = t2, shape = sh2, variant = v2)) &&
+            if v1 == AddMulVariant.MUL && v2 == AddMulVariant.MUL end => begin
+
+            if c1 isa Rat && c2 isa Rat
+                tmp = c1 // c2
+                c1 = tmp.num
+                c2 = tmp.den
+            end
+            n = Mul{T}(c1, d1, ; type = t1, shape = sh1)
+            d = Mul{T}(c2, d2, ; type = t2, shape = sh2)
+        end
+        _ => nothing
+    end
 
     _isone(d) && return Const{T}(n)
     _isone(-d) && return Const{T}(-n)

src/utils.jl

@@ -1,6 +1,6 @@
 using Base: ImmutableDict
 
-safe_isinteger(x::Number) = isinteger(x) && abs(x) < typemax(Int)
+safe_isinteger(@nospecialize(x::Number)) = isinteger(x) && abs(x) < typemax(Int)
 safe_isinteger(x) = false
 
 pow(x,y) = y==0 ? 1 : y<0 ? inv(x)^(-y) : x^y
@@ -29,13 +29,15 @@ isliteral(::Type{T}) where {T} = x -> x isa T
 is_literal_number(x) = isliteral(Number)(unwrap_const(x))
 
 # checking the type directly is faster than dynamic dispatch in type unstable code
-function _iszero(x)
+@cache function _iszero(x)::Bool
+    @nospecialize x
     x = unwrap_const(unwrap(x))
     x isa Number && return iszero(x)::Bool
     x isa Array && return iszero(x)::Bool
     return false
 end
-function _isone(x)
+@cache function _isone(x)::Bool
+    @nospecialize x
     x = unwrap_const(unwrap(x))
     x isa Number && return isone(x)::Bool
     x isa Array && return isone(x)::Bool

test/basics.jl

@@ -185,9 +185,8 @@ end
     @test isequal(w == 0, Term{SymReal}(==, [w, 0]; type = Bool))
 
     @syms x::Integer a::Integer
-    @eqtest x // 5 == (1 // 5) * x
+    @eqtest x // 5 == SymbolicUtils.Div{SymReal}(x, 5, false; type = Real)
     @eqtest (1//2 * x) / 5 == (1 // 10) * x
-    @eqtest x // Int16(5) == Rational{Int16}(1, 5) * x
     @eqtest 5 // x == 5 / x
     @eqtest x // a == x / a
 
@@ -1086,8 +1085,7 @@ end
     @test get_mul_coefficient((2x/3y).num) == 2
     @test get_mul_coefficient((2x/3y).den) == 3
     @test unwrap_const(2x/-3x) == -2//3
-    @test unwrap_const((2.5x/3x).num) == 2.5
-    @test unwrap_const((2.5x/3x).den) == 3
+    @test unwrap_const((2.5x/3x)) == 2.5/3
     @test unwrap_const(x/3x) == 1//3
     @test isequal(x / 1, x)
     @test isequal(x / -1, -x)

test/polyform.jl

@@ -64,7 +64,7 @@ end
     @test isequal(simplify_fractions(a), 7/expand(-(x-2)^2))
 
     # https://github.com/JuliaSymbolics/Symbolics.jl/issues/968
-    @eqtest simplify_fractions((x * y + (1//2) * x) / (2 * x)) == (1 + 2y) / 4
+    @eqtest simplify_fractions((x * y + (1//2) * x) / (2 * x)) == (1//2 + y) / 2
 end
 
 @testset "isone iszero" begin