Add tfunc support for copysign and ifelse

Fixes #18

Author: Keno

src/analysis/ipoincidence.jl

@@ -167,7 +167,7 @@ function _make_argument_lattice_elem(𝕃, which::Argument, @nospecialize(argt),
         return argt
     elseif Compiler.isprimitivetype(argt)
         inc = Incidence(add_variable!(which))
-        return argt === Float64 ? inc : Incidence(argt, inc.row, inc.eps)
+        return argt === Float64 ? inc : Incidence(argt, inc.row)
     elseif isa(argt, PartialStruct)
         return PartialStruct(𝕃, argt.typ, Any[make_argument_lattice_elem(𝕃, which, f, add_variable!, add_equation!, add_scope!) for f in argt.fields])
     elseif isabstracttype(argt) || ismutabletype(argt) || !isa(argt, DataType)

src/analysis/lattice.jl

@@ -430,7 +430,7 @@ function Compiler.tmeet(🥬::EqStructureLattice, @nospecialize(a), @nospecializ
         meet = Compiler.tmeet(Compiler.widenlattice(🥬), a.typ, b)
         meet == Union{} && return Union{}
         Base.issingletontype(meet) && return meet
-        return Incidence(meet, copy(a.row), copy(a.eps))
+        return Incidence(meet, copy(a.row))
     elseif isa(a, Eq)
         meet = Compiler.tmeet(Compiler.widenlattice(🥬), equation, b)
         meet == Union{} && return Union{}
@@ -514,15 +514,15 @@ function Compiler.tmerge(🥬::EqStructureLattice, @nospecialize(a), @nospeciali
                     row[i] = nonlinear
                 end
             end
-            return Incidence(merged_typ, row, union(a.eps, b.eps))
+            return Incidence(merged_typ, row)
         elseif isa(b, Const)
             # Const has no incidence taint
             typ = Compiler.tmerge(Compiler.widenlattice(🥬), a.typ, b)
             r = copy(a)
             for i in rowvals(r.row)
                 r.row[i] = nonlinear
             end
-            return Incidence(typ, r.row, copy(a.eps))
+            return Incidence(typ, r.row)
         else
             a = widenconst(a)
         end

src/analysis/refiner.jl

@@ -184,6 +184,22 @@ function tfunc(::Val{Core.Intrinsics.mul_float}, @nospecialize(a::Union{Const, T
     return Incidence(builtin_math_tfunc(Core.Intrinsics.mul_float, a.typ, b.typ), rrow)
 end
 
+function tfunc(::Val{Core.Intrinsics.copysign_float}, @nospecialize(a::Union{Const, Type{Float64}, Incidence}), @nospecialize(b::Union{Const, Type{Float64}, Incidence}))
+    if a === Float64 || b === Float64
+        return Float64
+    end
+    if isa(a, Const) && isa(b, Const)
+        return builtin_math_tfunc(Core.Intrinsics.copysign_float, a, b)
+    end
+    rrow = _zero_row()
+    arow = isa(a, Incidence) ? a.row : _ZERO_ROW
+    brow = isa(b, Incidence) ? b.row : _ZERO_ROW
+    for i in Iterators.flatten((rowvals(arow), rowvals(brow)))
+        rrow[i] = nonlinear
+    end
+    return Incidence(builtin_math_tfunc(Core.Intrinsics.copysign_float, widenconst(a), widenconst(b)), rrow)
+end
+
 function tfunc(::Val{Core.Intrinsics.div_float}, @nospecialize(a::Union{Const, Type{Float64}, Incidence}), @nospecialize(b::Union{Const, Type{Float64}, Incidence}))
     if isa(a, Const) && isa(b, Const)
         return builtin_math_tfunc(Core.Intrinsics.div_float, a, b)
@@ -258,18 +274,19 @@ is_any_incidence(@nospecialize args...) = any(@nospecialize(x)->isa(x, Incidence
         b = argtypes[2]
         if is_any_incidence(a, b)
             if (f == Core.Intrinsics.add_float || f == Core.Intrinsics.sub_float) ||
-                (f == Core.Intrinsics.mul_float || f == Core.Intrinsics.div_float)
+                (f == Core.Intrinsics.mul_float || f == Core.Intrinsics.div_float) ||
+                f == Core.Intrinsics.copysign_float
                 return tfunc(Val(f), a, b)
             elseif f == Core.Intrinsics.lt_float
-                r = tmerge(typeinf_lattice(interp), argtypes[1], argtypes[2])
+                r = Compiler.tmerge(Compiler.typeinf_lattice(interp), argtypes[1], argtypes[2])
                 @assert isa(r, Incidence)
-                return Incidence(Bool, r.row, r.eps)
+                return Incidence(Bool, r.row)
             elseif f === Core.getfield && isa(a, Incidence)
                 a = argtypes[1]
-                fT = getfield_tfunc(typeinf_lattice(interp), widenconst(a), argtypes[2])
+                fT = Compiler.getfield_tfunc(Compiler.typeinf_lattice(interp), widenconst(a), argtypes[2])
                 fT === Union{} && return Union{}
                 Base.issingletontype(fT) && return fT
-                return Incidence(fT, copy(a.row), copy(a.eps))
+                return Incidence(fT, copy(a.row))
             end
         end
     elseif length(argtypes) == 3
@@ -297,7 +314,19 @@ is_any_incidence(@nospecialize args...) = any(@nospecialize(x)->isa(x, Incidence
                         return c
                     end
                 end
-                # TODO: tmergea
+                rt = Compiler.tmerge(Compiler.typeinf_lattice(interp), b, c)
+                if isa(rt, Incidence)
+                    if isa(a, Incidence)
+                        rrow = copy(rt.row)
+                        for i in rowvals(a.row)
+                            rrow[i] = nonlinear
+                        end
+                        rt = Incidence(rt.typ, rrow)
+                    else
+                        rt = widenconst(rt)
+                    end
+                end
+                return rt
             end
         end
     end

src/transform/codegen/init_uncompress.jl

@@ -120,6 +120,7 @@ function gen_init_uncompress!(
 
             if is_known_invoke(stmt, variable, ir) || is_equation_call(stmt, ir)
                 display(ir)
+                @show stmt
                 error()
             elseif is_known_invoke(stmt, equation, ir)
                 # Equation - used, but only as an arg to equation call, which will all get

src/transform/tearing/schedule.jl

@@ -168,7 +168,9 @@ function schedule_nonlinear!(compact, param_vars, var_eq_matching, ir, val::Unio
         f = argextype(stmt.args[1], ir)
         @assert isa(f, Const)
         f = f.val
-        @assert f in (Core.Intrinsics.sub_float, Core.Intrinsics.add_float, Core.Intrinsics.mul_float)
+        @assert f in (Core.Intrinsics.sub_float, Core.Intrinsics.add_float,
+                      Core.Intrinsics.mul_float, Core.Intrinsics.copysign_float,
+                      Core.ifelse)
         call_is_linear = f in (Core.Intrinsics.sub_float, Core.Intrinsics.add_float)
     end
 

test/regression.jl

@@ -0,0 +1,24 @@
+module Regression
+
+using Test
+using DAECompiler
+using DAECompiler.Intrinsics
+using Sundials
+using SciMLBase
+using OrdinaryDiffEq
+
+const cf = Base.copysign_float
+const ief = Core.ifelse
+const -ᵢ = Core.Intrinsics.sub_float
+
+function tfb1()
+    x = continuous()
+    b = (x < 43200.)
+    v = ief(b, x, cf(0., x))
+    always!(ddt(x) -ᵢ v)
+end
+
+sol = solve(DAECProblem(tfb1, (1,) .=> 1.), IDA())
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(sol.t)))
+
+end

test/runtests.jl

@@ -1,3 +1,4 @@
 include("basic.jl")
 include("ipo.jl")
-include("ssrm.jl")
+include("ssrm.jl")
+include("regression.jl")