Fix varargs IPO test (#50)

Author: Keno

Manifest.toml

@@ -27,7 +27,10 @@ struct ADAnalyzer <: Compiler.AbstractInterpreter
 end
 
 Compiler.InferenceParams(interp::ADAnalyzer) = Compiler.InferenceParams()
-Compiler.OptimizationParams(interp::ADAnalyzer) = Compiler.OptimizationParams()
+Compiler.OptimizationParams(interp::ADAnalyzer) = Compiler.OptimizationParams(;
+        assume_fatal_throw = true,
+        compilesig_invokes = false
+    )
 Compiler.get_inference_world(interp::ADAnalyzer) = interp.world
 Compiler.get_inference_cache(interp::ADAnalyzer) = interp.inf_cache
 Compiler.cache_owner(::ADAnalyzer) = ADCache()
@@ -60,6 +63,8 @@ end
 struct AnalyzedSource
     ir::Compiler.IRCode
     inline_cost::Compiler.InlineCostType
+    nargs::UInt
+    isva::Bool
 end
 
 @override function Compiler.result_edges(interp::ADAnalyzer, caller::InferenceState)
@@ -70,7 +75,7 @@ end
 @override function Compiler.transform_result_for_cache(interp::ADAnalyzer, result::InferenceResult, edges::SimpleVector)
     ir = result.src.optresult.ir
     params = Compiler.OptimizationParams(interp)
-    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result))
+    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
 end
 
 @override function Compiler.transform_result_for_local_cache(interp::ADAnalyzer, result::InferenceResult)
@@ -79,7 +84,7 @@ end
     end
     ir = result.src.optresult.ir
     params = Compiler.OptimizationParams(interp)
-    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result))
+    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
 end
 
 function Compiler.retrieve_ir_for_inlining(ci::CodeInstance, result::AnalyzedSource)

src/analysis/ADAnalyzer.jl

@@ -155,12 +155,13 @@ function apply_linear_incidence(𝕃, ret::PartialStruct, caller::CallerMappingS
     return PartialStruct(𝕃, ret.typ, Any[apply_linear_incidence(𝕃, f, caller, mapping) for f in ret.fields])
 end
 
-function CalleeMapping(𝕃::Compiler.AbstractLattice, argtypes::Vector{Any}, callee_result::DAEIPOResult, template_argtypes)
+function CalleeMapping(𝕃::Compiler.AbstractLattice, argtypes::Vector{Any}, callee_ci::CodeInstance, callee_result::DAEIPOResult, template_argtypes)
     applied_scopes = Any[]
     coeffs = Vector{Any}(undef, length(callee_result.var_to_diff))
     eq_mapping = fill(0, length(callee_result.total_incidence))
 
-    process_template!(𝕃, coeffs, eq_mapping, applied_scopes, argtypes, template_argtypes)
+    va_argtypes = Compiler.va_process_argtypes(𝕃, argtypes, callee_ci.inferred.nargs, callee_ci.inferred.isva)
+    process_template!(𝕃, coeffs, eq_mapping, applied_scopes, va_argtypes, template_argtypes)
 
     return CalleeMapping(coeffs, eq_mapping, applied_scopes)
 end

src/analysis/ipoincidence.jl

@@ -57,9 +57,7 @@ Compiler.cache_owner(::StructuralRefiner) = StructureCache()
     end
 
     argtypes = Compiler.collect_argtypes(interp, stmt.args, Compiler.StatementState(nothing, false), irsv)[2:end]
-    m = Compiler.get_ci_mi(callee_codeinst).def
-    argtypes = Compiler.va_process_argtypes(Compiler.optimizer_lattice(interp), argtypes, UInt(m.nargs), m.isva)
-    mapping = CalleeMapping(Compiler.optimizer_lattice(interp), argtypes, callee_result, callee_codeinst.inferred.ir.argtypes)
+    mapping = CalleeMapping(Compiler.optimizer_lattice(interp), argtypes, callee_codeinst, callee_result, callee_codeinst.inferred.ir.argtypes)
     new_rt = apply_linear_incidence(Compiler.optimizer_lattice(interp), callee_result.extended_rt,
         CallerMappingState(callee_result, interp.var_to_diff, interp.varclassification, interp.varkinds, VarEqClassification[]), mapping)
 
@@ -87,7 +85,7 @@ function structural_inc_ddt(var_to_diff::DiffGraph, varclassification::Union{Vec
             push!(varclassification, varclassification[v])
         end
         if varkinds !== nothing
-            push!(varkinds, Intrinsics.Continuous)
+            push!(varkinds, varkinds[v])
         end
         add_edge!(var_to_diff, v, dv)
         return dv + 1

src/analysis/refiner.jl

@@ -80,6 +80,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
     empty!(ir.argtypes)
     (arg_replacements, new_argtypes, nexternalargvars) = flatten_arguments!(compact, old_argtypes, 0, ir.argtypes)
     for i = 1:nexternalargvars
+        # TODO: Need to handle different var kinds for IPO
         add_variable!(Argument(i))
     end
     argtypes = Any[Incidence(new_argtypes[i], i) for i = 1:nexternalargvars]
@@ -212,6 +213,12 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
             return UncompilableIPOResult(warnings, UnsupportedIRException("Saw invalid variable kind (`$kind`)  for variable $(var_num) (SSA $ssa)", ir))
         end
         varkinds[var_num] = kind.val
+        dv = var_num
+        while true
+            dv = var_to_diff[dv]
+            dv === nothing && break
+            varkinds[dv] = kind.val
+        end
 
         scope = argextype(inst.args[4], ir)
         if (!isa(scope, Const) || !isa(scope.val, Intrinsics.AbstractScope)) && !is_valid_partial_scope(scope)
@@ -330,7 +337,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
             end
 
             callee_argtypes = Any[argextype(stmt.args[i], compact) for i in 2:length(stmt.args)]
-            mapping = CalleeMapping(Compiler.optimizer_lattice(refiner), callee_argtypes, result, callee_codeinst.inferred.ir.argtypes)
+            mapping = CalleeMapping(Compiler.optimizer_lattice(refiner), callee_argtypes, callee_codeinst, result, callee_codeinst.inferred.ir.argtypes)
             inst[:info] = info = MappingInfo(info, result, mapping)
         end
 

src/analysis/structural.jl

@@ -96,6 +96,7 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
 
     all_states = Int[]
     for var = 1:length(result.var_to_diff)
+        varkind(state, var) == Intrinsics.Continuous || continue
         kind = classify_var(result.var_to_diff, key, var)
         kind == nothing && continue
         numstates[kind] += 1

src/transform/codegen/dae_factory.jl

@@ -85,6 +85,7 @@ function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
 
     all_states = Int[]
     for var = 1:length(result.var_to_diff)
+        varkind(state, var) == Intrinsics.Continuous || continue
         kind = classify_var(result.var_to_diff, key, var)
         kind == nothing && continue
         numstates[kind] += 1

src/transform/codegen/ode_factory.jl

@@ -25,7 +25,7 @@ function handle_contribution!(ir::Compiler.IRCode, inst::Compiler.Instruction, k
     replace_call!(ir, pos, Expr(:call, Base.setindex!, which, sum, slot))
 end
 
-function compute_slot_ranges(info::MappingInfo, callee_key, var_assignment, eq_assignment)
+function compute_slot_ranges(caller_state::TransformationState, info::MappingInfo, callee_key, var_assignment, eq_assignment)
     # Compute the ranges for this child's states in the parent range.
     # We rely upon earlier stages of the pipeline having put these adjacent to each other
     # and in order. We could just trust that, but because it's a little bit tricky, here
@@ -47,6 +47,7 @@ function compute_slot_ranges(info::MappingInfo, callee_key, var_assignment, eq_a
         caller_map = info.mapping.var_coeffs[callee_var]
         isa(caller_map, Const) && continue
         caller_var = only(rowvals(caller_map.row))-1
+        varkind(caller_state, caller_var) == Intrinsics.Continuous || continue
 
         callee_kind = classify_var(info.result.var_to_diff, callee_key, callee_var)
         callee_kind === nothing && continue
@@ -166,7 +167,7 @@ function rhs_finish!(
                 push!(stmt.args, in_vars)
 
                 # Ordering from tearing is (AssignedDiff, UnassignedDiff, Algebraic, Explicit)
-                slot_ranges = compute_slot_ranges(info, callee_key, var_assignment, eq_assignment)
+                slot_ranges = compute_slot_ranges(state, info, callee_key, var_assignment, eq_assignment)
                 for (arg, range) in zip(arg_range, slot_ranges)
                     push!(stmt.args, insert_node!(ir, SSAValue(i),
                         NewInstruction(inst;
@@ -183,6 +184,12 @@ function rhs_finish!(
                 error()
             elseif is_known_invoke(stmt, variable, ir)
                 varnum = idnum(ir.stmts.type[i])
+                kind = varkind(state, varnum)
+                if kind == Intrinsics.Epsilon
+                    replace_call!(ir, SSAValue(i), 0.)
+                    continue
+                end
+                @assert kind == Intrinsics.Continuous
 
                 assgn = var_assignment[varnum]
                 if assgn == nothing

src/transform/codegen/rhs.jl

@@ -101,7 +101,7 @@ function cache_dae_ci!(old_ci, src, debuginfo, abi, owner; rettype=Tuple)
     return daef_ci
 end
 
-function replace_call!(ir::Union{IRCode,IncrementalCompact}, idx::SSAValue, new_call::Expr)
+function replace_call!(ir::Union{IRCode,IncrementalCompact}, idx::SSAValue, @nospecialize(new_call))
     @assert !isa(ir[idx][:inst], PhiNode)
     ir[idx][:inst] = new_call
     ir[idx][:type] = Any

src/transform/common.jl

@@ -127,8 +127,18 @@ end
 dae_sol = solve(DAECProblem(sicm_vars(1.0), (1,) .=> 1.), IDA())
 ode_sol = solve(ODECProblem(sicm_vars(1.0), (1,) .=> 1.), Rodas5(autodiff=false))
 for sol in (dae_sol, ode_sol)
-    @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], 1. .+ sol.t))
+    @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[1, :], 1. .+ sol.t))
 end
 
+#= epsilon =#
+function simple_eps()
+    x = continuous()
+    always!(ddt(x) -ᵢ x +ᵢ epsilon())
+end
+dae_sol = solve(DAECProblem(simple_eps, (1,) .=> 1.), IDA())
+ode_sol = solve(ODECProblem(simple_eps, (1,) .=> 1.), Rodas5(autodiff=false))
+for sol in (dae_sol, ode_sol)
+    @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
+end
 
 end