Fix handling of functors, replace flatten_parameter!

Author: serenity4

src/analysis/flattening.jl

@@ -90,7 +90,7 @@ function flatten_argument!(state::FlatteningState, @nospecialize(argt))
         line = compact[Compiler.OldSSAValue(1)][:line]
         ssa = @insert_instruction_here(compact, line, settings, (:invoke)(nothing, InternalIntrinsics.external_equation)::Eq(eq))
         return ssa
-    elseif isa(argt, Type)
+    elseif argt <: Type
         return argt.parameters[1]
     elseif isabstracttype(argt) || ismutabletype(argt) || (!isa(argt, DataType) && !isa(argt, PartialStruct))
         line = compact[Compiler.OldSSAValue(1)][:line]
@@ -112,19 +112,19 @@ function flatten_argument!(state::FlatteningState, @nospecialize(argt))
     end
 end
 
-function flatten_arguments_for_callee!(compact::IncrementalCompact, map::ArgumentMap, argtypes, 𝕃, line, settings)
+function flatten_arguments_for_callee!(compact::IncrementalCompact, map::ArgumentMap, argtypes, args, line, settings, 𝕃 = Compiler.fallback_lattice)
     list = Any[]
     this = nothing
-    last_index = Int[]
+    last_index = CompositeIndex()
     for index in map.variables
         from = findfirst(j -> get(last_index, j, -1) !== index[j], eachindex(index))::Int
         for i in from:length(index)
             field = index[i]
             if i == 1
-                this = Argument(2 + field)
+                this = args[field]
             else
                 thistype = argextype(this, compact)
-                fieldtype = Compiler.getfield_tfunc(𝕃, Const(field))
+                fieldtype = Compiler.getfield_tfunc(𝕃, thistype, Const(field))
                 this = @insert_instruction_here(compact, line, settings, getfield(this, field)::fieldtype)
             end
         end
@@ -133,42 +133,6 @@ function flatten_arguments_for_callee!(compact::IncrementalCompact, map::Argumen
     return list
 end
 
-function _flatten_parameter!(𝕃, compact, argtypes, ntharg, line, settings)
-    list = Any[]
-    for (argn, argt) in enumerate(argtypes)
-        if isa(argt, Const)
-            continue
-        elseif Base.issingletontype(argt)
-            continue
-        elseif Base.isprimitivetype(argt) || isa(argt, Incidence)
-            push!(list, ntharg(argn))
-        elseif argt === equation || isa(argt, Eq)
-            continue
-        elseif isa(argt, Type) && argt <: Intrinsics.AbstractScope
-            continue
-        elseif isabstracttype(argt) || ismutabletype(argt) || (!isa(argt, DataType) && !isa(argt, PartialStruct))
-            continue
-        else
-            if !isa(argt, PartialStruct) && Base.datatype_fieldcount(argt) === nothing
-                continue
-            end
-            this = ntharg(argn)
-            nthfield(i) = @insert_instruction_here(compact, line, settings, getfield(this, i)::Compiler.getfield_tfunc(𝕃, argextype(this, compact), Const(i)))
-            if isa(argt, PartialStruct)
-                fields = _flatten_parameter!(𝕃, compact, argt.fields, nthfield, line, settings)
-            else
-                fields = _flatten_parameter!(𝕃, compact, fieldtypes(argt), nthfield, line, settings)
-            end
-            append!(list, fields)
-        end
-    end
-    return list
-end
-
-function flatten_parameter!(𝕃, compact, argtypes, ntharg, line, settings)
-    return @insert_instruction_here(compact, line, settings, tuple(_flatten_parameter!(𝕃, compact, argtypes, ntharg, line, settings)...)::Tuple)
-end
-
 remove_variable_and_equation_annotations(argtypes) = Any[widenconst(T) for T in argtypes]
 
 function annotate_variables_and_equations(argtypes::Vector{Any}, map::ArgumentMap)

src/analysis/structural.jl

@@ -376,7 +376,11 @@ function _structural_analysis!(ci::CodeInstance, world::UInt, settings::Settings
             # Rewrite to flattened ABI
             compact[SSAValue(i)] = nothing
             compact.result_idx -= 1
-            new_args = _flatten_parameter!(Compiler.optimizer_lattice(refiner), compact, callee_codeinst.inferred.ir.argtypes, arg->stmt.args[arg+1], line, settings)
+            callee_argtypes = callee_codeinst.inferred.ir.argtypes
+            callee_argmap = ArgumentMap(callee_argtypes)
+            args = @view(stmt.args[2:end])
+            𝕃 = Compiler.optimizer_lattice(refiner)
+            new_args = flatten_arguments_for_callee!(compact, callee_argmap, callee_argtypes, args, line, settings, 𝕃)
             new_call = insert_instruction_here!(compact, settings, @__SOURCE__,
             NewInstruction(Expr(:invoke, (StructuralSSARef(compact.result_idx), callee_codeinst), new_args...), stmtype, info, line, stmtflags))
             compact.ssa_rename[compact.idx - 1] = new_call

src/transform/codegen/dae_factory.jl

@@ -147,7 +147,7 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Unio
     argt = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Float64}
     sicm = ()
     if settings.skip_optimizations
-        daef_ci = rhs_finish_noopt!(state, ci, key, world, settings)
+        daef_ci = rhs_finish_noopt!(state, ci, key, world, settings; opaque_closure = true)
         oc = sciml_to_internal_abi_noopt!(copy(ci.inferred.ir), state, daef_ci, settings)
     else
         # TODO: We should not have to recompute this here
@@ -164,8 +164,12 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Unio
             @assert sicm_ci !== nothing
 
             line = result.ir[SSAValue(1)][:line]
-            param_list = flatten_parameter!(Compiler.fallback_lattice, compact, ci.inferred.ir.argtypes[1:end], argn->Argument(2+argn), line, settings)
-            sicm = @insert_instruction_here(compact, line, settings, invoke(param_list, sicm_ci)::Tuple)
+            callee_argtypes = ci.inferred.ir.argtypes
+            callee_argmap = ArgumentMap(callee_argtypes)
+            args = Argument.(2 .+ eachindex(callee_argtypes))
+            new_args = flatten_arguments_for_callee!(compact, callee_argmap, callee_argtypes, args, line, settings)
+            list = @insert_instruction_here(compact, line, settings, tuple(new_args...)::Tuple)
+            sicm = @insert_instruction_here(compact, line, settings, invoke(list, sicm_ci)::Tuple)
         end
 
         daef_ci = rhs_finish!(state, ci, key, world, settings, 1)

src/transform/codegen/init_factory.jl

@@ -27,7 +27,11 @@ function init_uncompress_gen!(compact::Compiler.IncrementalCompact, result::DAEI
         @assert sicm_ci !== nothing
 
         line = result.ir[SSAValue(1)][:line]
-        param_list = flatten_parameter!(Compiler.fallback_lattice, compact, ci.inferred.ir.argtypes[1:end], argn->Argument(2+argn), line, settings)
+        callee_argtypes = ci.inferred.ir.argtypes
+        callee_argmap = ArgumentMap(callee_argtypes)
+        args = Argument.(2 .+ eachindex(callee_argtypes))
+        new_args = flatten_arguments_for_callee!(compact, callee_argmap, callee_argtypes, args, line, settings)
+        param_list = @insert_instruction_here(compact, line, settings, tuple(new_args...)::Tuple)
         sicm = @insert_instruction_here(compact, line, settings, invoke(param_list, sicm_ci)::Tuple)
     else
         sicm = ()

src/transform/codegen/ode_factory.jl

@@ -70,8 +70,12 @@ function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
         @assert sicm_ci !== nothing
 
         line = result.ir[SSAValue(1)][:line]
-        param_list = flatten_parameter!(Compiler.fallback_lattice, returned_ic, ci.inferred.ir.argtypes[1:end], argn->Argument(2+argn), line, settings)
-        sicm_state = @insert_instruction_here(returned_ic, line, settings, (:call)(invoke, param_list, sicm_ci)::Tuple)
+        callee_argtypes = ci.inferred.ir.argtypes
+        callee_argmap = ArgumentMap(callee_argtypes)
+        args = Argument.(2 .+ eachindex(callee_argtypes))
+        new_args = flatten_arguments_for_callee!(returned_ic, callee_argmap, callee_argtypes, args, line, settings)
+        param_list = @insert_instruction_here(returned_ic, line, settings, tuple(new_args...)::Tuple)
+        sicm_state = @insert_instruction_here(returned_ic, line, settings, invoke(param_list, sicm_ci)::Tuple)
     else
         sicm_state = ()
     end

src/transform/reconstruct.jl

@@ -94,7 +94,8 @@ derivative may differ between the unoptimized and optimized versions.
 function compute_residual_vectors(f, u, du; t = 1.0, mode=DAE, world=Base.tls_world_age())
     @assert mode === DAE # TODO: support ODEs
     settings = Settings(; mode, insert_stmt_debuginfo = true)
-    ci = _code_ad_by_type(Tuple{typeof(f)}; world)
+    tt = Base.signature_type(f, ())
+    ci = _code_ad_by_type(tt; world)
     result = @code_structure result=true mode=mode world=world f()
     structure = make_structure_from_ipo(result)
     state = TransformationState(result, structure)

src/transform/unoptimized.jl

@@ -13,7 +13,8 @@ function rhs_finish_noopt!(
     world::UInt,
     settings::Settings,
     equation_to_residual_mapping = 1:length(state.structure.eq_to_diff),
-    variable_to_state_mapping = map_variables_to_states(state))
+    variable_to_state_mapping = map_variables_to_states(state);
+    opaque_closure)
 
     (; result, structure) = state
     result_ci = find_matching_ci(ci -> ci.owner === key, ci.def, world)
@@ -23,12 +24,14 @@ function rhs_finish_noopt!(
 
     ir = copy(result.ir)
     src = ci.inferred::AnalyzedSource
-    argrange = 2:src.nargs
-    slotnames = Symbol[:captures]
-    argtypes = Any[Tuple]
+    argrange = 1:src.nargs
     # Original arguments.
-    append!(slotnames, src.slotnames[argrange])
-    append!(argtypes, remove_variable_and_equation_annotations(ir.argtypes[argrange]))
+    slotnames = src.slotnames[argrange]
+    argtypes = remove_variable_and_equation_annotations(ir.argtypes)
+    if opaque_closure
+        slotnames[1] = :captures
+        argtypes[1] = Tuple
+    end
     # Additional ABI arguments.
     push!(slotnames, :out, :du, :u, :residuals, :states, :t)
     push!(argtypes, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Int}, Vector{Int}, Float64)
@@ -88,7 +91,7 @@ function rhs_finish_noopt!(
             inst[:stmt] = 0.0
         elseif isexpr(stmt, :invoke)
             info = inst[:info]::MappingInfo
-            callee_ci, callee_f, original_args = stmt.args[1]::CodeInstance, stmt.args[2], @view stmt.args[3:end]
+            callee_ci, args = stmt.args[1]::CodeInstance, @view stmt.args[2:end]
             callee_result = structural_analysis!(callee_ci, world, settings)
             callee_structure = make_structure_from_ipo(callee_result)
             callee_state = TransformationState(callee_result, callee_structure)
@@ -102,9 +105,8 @@ function rhs_finish_noopt!(
             end
             callee_states = [get(variable_to_state_mapping, i, -1) for i in caller_variables]
 
-            callee_daef_ci = rhs_finish_noopt!(callee_state, callee_ci, UnoptimizedKey(), world, settings, callee_residuals, callee_states)
-            call = @insert_instruction_here(compact, line, settings, (:invoke)(callee_daef_ci, callee_f,
-                original_args...,
+            callee_daef_ci = rhs_finish_noopt!(callee_state, callee_ci, UnoptimizedKey(), world, settings, callee_residuals, callee_states; opaque_closure = false)
+            call = @insert_instruction_here(compact, line, settings, (:invoke)(callee_daef_ci, args...,
                 out,
                 du,
                 u,

test/validation.jl

@@ -85,6 +85,19 @@ function equation_with_callable!()
     always!(ddt(callable()) - 3.0)
 end
 
+@noinline apply_equation!(lhs, rhs) = always!(lhs - rhs)
+function nonlinear_argument!()
+    x = continuous()
+    apply_equation!(ddt(x), sin(x))
+end
+
+struct WithParameter{N} end
+@noinline (::WithParameter{N})(eq, x) where {N} = eq(ddt(x) - N)
+function callable_with_type_parameter!()
+    eq, x = new_equation_and_variable()
+    WithParameter{3}()(eq, x)
+end
+
 @testset "Validation" begin
     refresh() # TODO: remove before merge
 
@@ -132,6 +145,12 @@ end
     @test all(>(0), residuals)
     @test residuals ≈ expanded_residuals
 
+    u = [2.0]
+    du = [4.0]
+    residuals, expanded_residuals = compute_residual_vectors(nonlinear_argument!, u, du)
+    @test residuals ≈ du .- sin.(u)
+    @test residuals ≈ expanded_residuals
+
     u = [0.0]
     du = [2.0]
     residuals, expanded_residuals = compute_residual_vectors(external_equation!, u, du)
@@ -161,4 +180,11 @@ end
     residuals, expanded_residuals = compute_residual_vectors(equation_with_callable!, u, du)
     @test residuals ≈ [1.0]
     @test residuals ≈ expanded_residuals
+
+    u = [2.0]
+    du = [3.0]
+    refresh(); residuals, expanded_residuals = compute_residual_vectors(callable_with_type_parameter!, u, du)
+    @test residuals ≈ [0.0]
+    # XXX: Residuals from the optimized pipeline are wrong (3.0 instead of 0.0)
+    @test_broken residuals ≈ expanded_residuals
 end;