Implement rewrite validator (#55)

* Misc refactors, insert internal stmt debuginfo

* Add slotnames info

* Refactor naming of locals for ODE codegen

* Add comment

* Apply DebugInfo at top-level

* Remove `insert_ssa_debuginfo` setting

* Marshall settings everywhere

* Add maybe_insert_debuginfo

* Disable crashing tests

* Propagate source information for `insert_node_here!(...)`

* Add source provenance to DAE/Init codegen

* Use `__SOURCE__` macro for `replace_call!`

* Update `replace_call!` with source information in index_lowering.jl

* Default `insert_stmt_debuginfo` to `false` for reflection tools

* Remove accidental code inclusion

* Don't seek previous codeloc

* Remove unused code

* Reenable IPO tests

* [DO NOT MERGE] Temporarily dev ConstructionBase for CI

* Work around  `invokelatest` issue

* Only wrap `string` call in `try`/`catch`

* Add `insert_ssa_debuginfo` setting

* Minor fix

* Refactor `insert_node_here` macro

* Use insert_instruction! in more places

* Rename `insert_node_here` macro to `insert_instruction`

* `insert_instruction` -> `insert_instruction_here`

Not to confuse it with the nonlocal insertion

* Fixes

* Undev ConstructionBase (and update)

* [WIP] Add unoptimized compilation path

* Use UnoptimizedKey singleton for caching unoptimized RHS

* Refactor settings construction for DAE/ODE problems

* Write `expand_residuals`

* Use :call form for macrocall

* Use :call for added macrocalls

* Fix introduced codegen bug

* [WIP] Add basic infrastructure for testing

* Reconstruct variables from states

* Separate optimized and unoptimized structural caches

* Add incidence-copying constructor for TransformationState

* Minor refactor `rhs_ir_finish!`

* Skip flattening in unoptimized context, WIP on IPO support

* Correctly compute state/equation indices for IPO

This includes a refactor that considers mapping to states,
instead of individual mappings into u/du vectors

* Adjust residual sign based on solved variable coefficient

* Rewrite and support flattening, improve external equation support

* Fix flattening bug

* Fix another flattening bug

* Rename testing function

* Support type constructors for IPO

* Split benchmark definition and test

* Remove process_template!

* Fix handling of functors, replace flatten_parameter!

* Fix bug in residual expansion

* Mark `apply_linear_incidence` as mutating

* Correctly handle nonlinear replacements

* Add (broken) validation test for thermalfluid benchmark

* Wrap test files in modules

* Remove unnecessary import

* Remove temp package from project

* Remove `refresh()` in tests

Author: serenity4

Manifest.toml

@@ -0,0 +1,10 @@
+using Test
+using SciMLBase, Sundials
+
+include("thermalfluid.jl")
+
+Benchmark{3}()()
+@test isa(code_lowered(DAECompiler.factory, Tuple{Val{DAECompiler.Settings(mode=DAECompiler.DAENoInit)}, Benchmark{3}})[1], Core.CodeInfo)
+let sol = solve(DAECProblem(Benchmark{3}(), [1:9;] .=> 0.), IDA())
+    @test_broken sol.retcode == ReturnCode.Success
+end

benchmark/main.jl

@@ -2,8 +2,6 @@ using DAECompiler
 using DAECompiler.Intrinsics
 using Polynomials: fit
 using XSteam: my_pT, rho_pT, Cp_pT, tc_pT
-using SciMLBase, Sundials
-using Test
 
 struct Sink; end
 function (::Sink)(inlet)
@@ -282,9 +280,3 @@ function (::Benchmark{N})() where {N}
     PreinsulatedPipe{N}()(in[2], out[1])
     Sink()(out[2])
 end
-
-Benchmark{3}()()
-@test isa(code_lowered(DAECompiler.factory, Tuple{Val{DAECompiler.Settings(mode=DAECompiler.DAENoInit)}, Benchmark{3}})[1], Core.CodeInfo)
-let sol = solve(DAECProblem(Benchmark{3}(), [1:9;] .=> 0.), IDA())
-    @test_broken sol.retcode == ReturnCode.Success
-end

benchmark/thermalfluid.jl

@@ -6,7 +6,7 @@ module DAECompiler
     using Diffractor
     using OrderedCollections
     using Compiler
-    using Compiler: IRCode, IncrementalCompact, DebugInfoStream, NewInstruction, argextype, singleton_type, isexpr, widenconst
+    using Compiler: AbstractLattice, IRCode, IncrementalCompact, DebugInfoStream, NewInstruction, argextype, singleton_type, isexpr, widenconst
     using Core.IR
     using SciMLBase
     using AutoHashEquals
@@ -16,18 +16,21 @@ module DAECompiler
     include("settings.jl")
     include("utils.jl")
     include("intrinsics.jl")
+    include("reflection.jl")
     include("analysis/utils.jl")
     include("analysis/lattice.jl")
     include("analysis/ADAnalyzer.jl")
     include("analysis/scopes.jl")
+    include("analysis/flattening.jl")
     include("analysis/cache.jl")
     include("analysis/refiner.jl")
     include("analysis/ipoincidence.jl")
     include("analysis/structural.jl")
-    include("analysis/flattening.jl")
     include("transform/state_selection.jl")
     include("transform/common.jl")
     include("transform/runtime.jl")
+    include("transform/unoptimized.jl")
+    include("transform/reconstruct.jl")
     include("transform/tearing/schedule.jl")
     include("transform/codegen/dae_factory.jl")
     include("transform/codegen/ode_factory.jl")
@@ -40,5 +43,4 @@ module DAECompiler
     include("analysis/consistency.jl")
     include("interface.jl")
     include("problem_interface.jl")
-    include("reflection.jl")
 end

src/DAECompiler.jl

@@ -62,6 +62,7 @@ end
 
 struct AnalyzedSource
     ir::Compiler.IRCode
+    slotnames::Vector{Any}
     inline_cost::Compiler.InlineCostType
     nargs::UInt
     isva::Bool
@@ -72,19 +73,26 @@ end
     Core.svec(edges..., interp.edges...)
 end
 
+function get_slotnames(def::Method)
+    names = split(def.slot_syms, '\0')
+    return map(Symbol, names)
+end
+
 @override function Compiler.transform_result_for_cache(interp::ADAnalyzer, result::InferenceResult, edges::SimpleVector)
     ir = result.src.optresult.ir
+    slotnames = get_slotnames(result.linfo.def)
     params = Compiler.OptimizationParams(interp)
-    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
+    return AnalyzedSource(ir, slotnames, 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)
     if Compiler.result_is_constabi(interp, result)
         return nothing
     end
     ir = result.src.optresult.ir
+    slotnames = get_slotnames(result.linfo.def)
     params = Compiler.OptimizationParams(interp)
-    return AnalyzedSource(ir, Compiler.compute_inlining_cost(interp, result), result.src.src.nargs, result.src.src.isva)
+    return AnalyzedSource(ir, slotnames, 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

@@ -21,6 +21,19 @@ end
     var_schedule::Vector{Pair{BitSet, BitSet}}
 end
 
+struct VariableReplacement
+    replaced::Incidence
+    by::Int
+    equation::Int
+end
+
+mutable struct NonlinearReplacementMap
+    const variables::Vector{VariableReplacement}
+    variable_counter::Int
+    equation_counter::Int
+end
+NonlinearReplacementMap() = NonlinearReplacementMap(VariableReplacement[], 0, 0)
+
 """
     StructuralSSARef
 
@@ -35,6 +48,7 @@ struct DAEIPOResult
     opaque_eligible::Bool
     extended_rt::Any
     argtypes
+    argmap::ArgumentMap
     nexternalargvars::Int # total vars is length(var_to_diff)
     nsysmscopes::Int
     nexternaleqs::Int
@@ -44,6 +58,7 @@ struct DAEIPOResult
     total_incidence::Vector{Any}
     eqclassification::Vector{VarEqClassification}
     eq_callee_mapping::Vector{Union{Nothing, Vector{Pair{StructuralSSARef, Int}}}}
+    replacement_map::NonlinearReplacementMap
     names::OrderedDict{Any, ScopeDictEntry} # TODO: OrderedIdDict
     varkinds::Vector{Union{Intrinsics.VarKind, Nothing}}
     eqkinds::Vector{Union{Intrinsics.EqKind, Nothing}}