Enzyme reverse mode fails in the inline linear solve's task-local cache lookup (`inline_linear_sccs = true`)

Aperta
  • Qualcosa non va?
#183 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
48/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
julia
Ambito
backend-api-design

Direzione di ricerca

Start with get_inline_linsolve_cache and numeric_ldiv! in ModelingToolkitTearing/src/reassemble.jl, especially the task-local storage lookup described in the report. Trace how remake rebuilds the initialization problem and invokes the observed function during Enzyme's augmented forward pass. Reproduce with inline_linear_sccs = true and the reported remake plus solve gradient; done when reverse-mode differentiation succeeds through the lookup.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

With DefaultReassembleAlgorithm(; inline_linear_sccs = true), the generated code solves linear SCCs through numeric_ldiv!, which looks up a LinearCache in task_local_storage() (get_inline_linsolve_cache). An Enzyme reverse-mode gradient through remake + solve of such a problem fails in that lookup, during Enzyme's augmented forward pass:

ERROR: LoadError: EnzymeRuntimeException: Enzyme execution failed.
Enzyme: Not yet implemented constant table in jl_eqtable_get   %"d::IdDict.ht" = load atomic ptr ...
Stacktrace:
   [1] get!
   [2] augmented_julia_get__87649wrap
   ...
   [7] runtime_generic_augfwd(...)
   [8] get_inline_linsolve_cache
   [9] numeric_ldiv!
  [10] safe_ldiv
   ...
  [23] ObservedWrapper
   ...
  [27] ReconstructInitializeprob
   ...
  [35] _remake_initialization_data_impl
   ...
  [47] remake
  [48] loss

The failing call is the observed function that remake evaluates to rebuild the initialization problem, so the failure happens before solve starts. The right-hand side calls the same inline solve; whether it fails the same way once remake gets through is not known.

Full stacktrace (inline, Tsit5(); repro.jl:72 is the Enzyme gradient call, in the script with the version-printing loop) 
ERROR: LoadError: EnzymeRuntimeException: Enzyme execution failed.
Enzyme: Not yet implemented constant table in jl_eqtable_get   %"d::IdDict.ht" = load atomic ptr addrspace(10), ptr addrspace(11) %1 unordered, align 8, !dbg !26, !tbaa !30, !alias.scope !34, !noalias !37, !nonnull !0, !dereferenceable !42, !dereferenceable_or_null !43, !align !43, !enzyme_type !44, !enzymejl_byref_MUT_REF !0   %3 = call nonnull "enzyme_type"="{[-1]:Pointer}" ptr addrspace(10) @ijl_eqtable_get(ptr addrspace(10) nonnull %"d::IdDict.ht", ptr addrspace(10) nonnull %"key::Symbol", ptr addrspace(10) %2) #13, !dbg !29 result: (false, nothing) (false, nothing, nothing) dict: (false, nothing) (true, Memory{Any}, GPUCompiler.MUT_REF) key (false, nothing) (true, Symbol, GPUCompiler.MUT_REF) dflt (true, :__c782dbf1cf4d6a2e5e3865d7e95634f2e09b5902__) (true, Symbol, GPUCompiler.BITS_REF)

Stacktrace:
   [1] get!
     @ ./iddict.jl:169 [inlined]
   [2] augmented_julia_get__89710wrap
     @ ./iddict.jl:0
   [3] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:8175 [inlined]
   [4] enzyme_call
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7588 [inlined]
   [5] AugmentedForwardThunk
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7534 [inlined]
   [6] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:467 [inlined]
   [7] runtime_generic_augfwd(::Type{Val{(false, false, false, false)}}, ::Val{true}, ::Val{false}, ::Val{1}, ::Val{(true, true, true, true)}, ::Val{@NamedTuple{var"1", var"2", var"3"}}, ::typeof(get!), ::Nothing, ::Type{Dict{Tuple{DataType, Tuple{Int64, Int64}}, Any}}, ::Nothing, ::IdDict{Any, Any}, ::IdDict{Any, Any}, ::Symbol, ::Nothing)
     @ Enzyme.Compiler ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:600
   [8] get_inline_linsolve_cache
     @ ~/.julia/packages/ModelingToolkitTearing/5tRoY/src/reassemble.jl:760 [inlined]
   [9] numeric_ldiv!
     @ ~/.julia/packages/ModelingToolkitTearing/5tRoY/src/reassemble.jl:737
  [10] safe_ldiv
     @ ~/.julia/packages/ModelingToolkitTearing/5tRoY/src/reassemble.jl:715 [inlined]
  [11] macro expansion
     @ ~/.julia/packages/RuntimeGeneratedFunctions/shJFB/src/RuntimeGeneratedFunctions.jl:243 [inlined]
  [12] macro expansion
     @ ./none:0 [inlined]
  [13] generated_callfunc
     @ ./none:0
  [14] RuntimeGeneratedFunction
     @ ~/.julia/packages/RuntimeGeneratedFunctions/shJFB/src/RuntimeGeneratedFunctions.jl:230 [inlined]
  [15] GeneratedFunctionWrapper
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/codegen_utils.jl:1216 [inlined]
  [16] GeneratedFunctionWrapper
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/codegen_utils.jl:0 [inlined]
  [17] augmented_julia_GeneratedFunctionWrapper_77772_inner_46wrap
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/codegen_utils.jl:0
  [18] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:8175 [inlined]
  [19] enzyme_call
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7588 [inlined]
  [20] AugmentedForwardThunk
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7534 [inlined]
  [21] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:467 [inlined]
  [22] runtime_generic_augfwd(::Type{Val{(false, false, true, true)}}, ::Val{true}, ::Val{false}, ::Val{1}, ::Val{(true, true, true, true)}, ::Val{@NamedTuple{var"1", var"2", var"3"}}, ::ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :__argₛᵧₘ8912177291400142833), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xe1142bf0, 0x7a3fecc5, 0xd4ebc7be, 0x12b0aa24, 0x0708801d), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x1, :x2, :x3, :x4), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0x31c3e887, 0x0dcece7a, 0x0f8d7ecc, 0xc94f4953, 0xe209970a), Nothing}}, ::ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :__argₛᵧₘ8912177291400142833), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xe1142bf0, 0x7a3fecc5, 0xd4ebc7be, 0x12b0aa24, 0x0708801d), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x1, :x2, :x3, :x4), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0x31c3e887, 0x0dcece7a, 0x0f8d7ecc, 0xc94f4953, 0xe209970a), Nothing}}, ::Vector{Float64}, ::Vector{Float64}, ::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, ::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, ::Float64, ::Base.RefValue{Float64})
     @ Enzyme.Compiler ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:600
  [23] ObservedWrapper
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/problem_utils.jl:819
  [24] call_composed
     @ ./operators.jl:1111 [inlined]
  [25] call_composed
     @ ./operators.jl:1110 [inlined]
  [26] ComposedFunction
     @ ./operators.jl:1107 [inlined]
  [27] ReconstructInitializeprob
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/problem_utils.jl:1484 [inlined]
  [28] ReconstructInitializeprob
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/problem_utils.jl:0 [inlined]
  [29] augmented_julia_ReconstructInitializeprob_76703_inner_29wrap
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/problem_utils.jl:0
  [30] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:8175 [inlined]
  [31] enzyme_call
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7588 [inlined]
  [32] AugmentedForwardThunk
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7534 [inlined]
  [33] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:467 [inlined]
  [34] runtime_generic_augfwd(::Type{Val{(true, true, true)}}, ::Val{true}, ::Val{false}, ::Val{1}, ::Val{(true, true, true)}, ::Val{@NamedTuple{var"1", var"2", var"3"}}, ::ModelingToolkitBase.ReconstructInitializeprob{ModelingToolkitBase.MTKParametersReconstructor{ComposedFunction{ModelingToolkitBase.PConstructorApplicator{typeof(identity)}, ModelingToolkitBase.CopyParamsByTemplate{true, Vector{Union{ModelingToolkitBase.IndepVarTemplate, ModelingToolkitBase.ParameterIndex{SciMLStructures.Initials, UnitRange{Int64}}, ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, UnitRange{Int64}}}}, 1, Nothing}}, Returns{StaticArraysCore.SVector{0, Float64}}, Returns{Tuple{}}, Returns{Tuple{}}, Returns{Tuple{}}}, ComposedFunction{typeof(identity), ModelingToolkitBase.ObservedWrapper{true, Any}}}, ::ModelingToolkitBase.ReconstructInitializeprob{ModelingToolkitBase.MTKParametersReconstructor{ComposedFunction{ModelingToolkitBase.PConstructorApplicator{typeof(identity)}, ModelingToolkitBase.CopyParamsByTemplate{true, Vector{Union{ModelingToolkitBase.IndepVarTemplate, ModelingToolkitBase.ParameterIndex{SciMLStructures.Initials, UnitRange{Int64}}, ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, UnitRange{Int64}}}}, 1, Nothing}}, Returns{StaticArraysCore.SVector{0, Float64}}, Returns{Tuple{}}, Returns{Tuple{}}, Returns{Tuple{}}}, ComposedFunction{typeof(identity), ModelingToolkitBase.ObservedWrapper{true, Any}}}, ::SymbolicIndexingInterface.ProblemState{Vector{Float64}, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, Float64, Nothing}, ::Base.RefValue{SymbolicIndexingInterface.ProblemState{Vector{Float64}, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, Float64, Nothing}}, ::SCCNonlinearProblem{Vector{Float64}, false, Tuple{LinearProblem{Vector{Float64}, true, Matrix{Float64}, Vector{Float64}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, SciMLBase.SymbolicLinearInterface{SciMLBase.UpdateABWrapper{FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Matrix{Float64}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{LinearAlgebra.Diagonal{Float64, Vector{Float64}}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{BandedMatrices.BandedMatrix{Float64, Matrix{Float64}, Base.OneTo{Int64}}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}}, FunctionWrappersWrappers.AllowNonIsBits, FunctionWrappersWrappers.SingleCacheStorage}, FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{Float64}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}}, FunctionWrappersWrappers.AllowNonIsBits, FunctionWrappersWrappers.SingleCacheStorage}}, System, ModelingToolkitBase.ObservedFunctionCache{System, Nothing}, Nothing}, @Kwargs{verbose::SciMLLogging.Minimal}}}, Tuple{Returns{Nothing}}, NonlinearFunction{false, SciMLBase.AutoSpecialize, Returns{Nothing}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, System, Nothing, Nothing}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, Val{true}}, ::SCCNonlinearProblem{Vector{Float64}, false, Tuple{LinearProblem{Vector{Float64}, true, Matrix{Float64}, Vector{Float64}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, SciMLBase.SymbolicLinearInterface{SciMLBase.UpdateABWrapper{FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Matrix{Float64}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{LinearAlgebra.Diagonal{Float64, Vector{Float64}}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{BandedMatrices.BandedMatrix{Float64, Matrix{Float64}, Base.OneTo{Int64}}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}}, FunctionWrappersWrappers.AllowNonIsBits, FunctionWrappersWrappers.SingleCacheStorage}, FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{Float64}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}}}}, FunctionWrappersWrappers.AllowNonIsBits, FunctionWrappersWrappers.SingleCacheStorage}}, System, ModelingToolkitBase.ObservedFunctionCache{System, Nothing}, Nothing}, @Kwargs{verbose::SciMLLogging.Minimal}}}, Tuple{Returns{Nothing}}, NonlinearFunction{false, SciMLBase.AutoSpecialize, Returns{Nothing}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, System, Nothing, Nothing}, MTKParameters{Vector{Float64}, StaticArraysCore.SVector{0, Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, Val{true}})
     @ Enzyme.Compiler ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:600
  [35] _remake_initialization_data_impl
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/nonlinear/initializesystem.jl:706
  [36] remake_initialization_data
     @ ~/.julia/packages/ModelingToolkitBase/mpZiU/src/systems/nonlinear/initializesystem.jl:671 [inlined]
  [37] remake_initialization_data
     @ ~/.julia/packages/SciMLBase/Zzjmy/src/remake.jl:887 [inlined]
  [38] #_remake_odeproblem#963
     @ ~/.julia/packages/SciMLBase/Zzjmy/src/remake.jl:663 [inlined]
  [39] augmented_julia___remake_odeproblem_963_73956_inner_3wrap
     @ ~/.julia/packages/SciMLBase/Zzjmy/src/remake.jl:0
  [40] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:8175 [inlined]
  [41] enzyme_call
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7588 [inlined]
  [42] AugmentedForwardThunk
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7534 [inlined]
  [43] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:467 [inlined]
  [44] runtime_generic_augfwd(::Type{Val{(false, false, false, false, true, false, false, false, false, false, false, false, false)}}, ::Val{true}, ::Val{false}, ::Val{1}, ::Val{(true, true, true, true, true, true, true, true, true, true, true, true, true)}, ::Val{@NamedTuple{var"1", var"2", var"3"}}, ::SciMLBase.var"##_remake_odeproblem#963", ::Nothing, ::Missing, ::Nothing, ::Vector{Float64}, ::Nothing, ::Missing, ::Nothing, ::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, ::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, ::Missing, ::Nothing, ::Bool, ::Nothing, ::Type{Val{true}}, ::Nothing, ::Bool, ::Nothing, ::Nothing, ::Nothing, ::@Kwargs{}, ::Nothing, ::typeof(SciMLBase._remake_odeproblem), ::Nothing, ::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, ODEFunction{true, SciMLBase.AutoDespecialize, ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:___mtkunknowns___, :___mtkparameters___, :__argₛᵧₘ8912177291400142833), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0x8b376b31, 0x5e9e5a0d, 0xe002801d, 0xcb558804, 0x62e6ffb2), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__argₛᵧₘ14743185283305486891, :___mtkunknowns___, :___mtkparameters___, :__argₛᵧₘ8912177291400142833), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0x262a6e1c, 0xd91a41ee, 0xc3251555, 0x9af809af, 0x7b208f3b), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, ModelingToolkitBase.ObservedFunctionCache{System, Nothing}, Nothing, System, Union{Nothing, SciMLBase.OverrideInitData}, Union{Nothing, SciMLBase.ODENLStepData}}, @Kwargs{}, SciMLBase.StandardODEProblem}, ::Nothing)
     @ Enzyme.Compiler ~/.julia/packages/Enzyme/wlmJK/src/rules/jitrules.jl:600
  [45] _remake_odeproblem
     @ ~/.julia/packages/SciMLBase/Zzjmy/src/remake.jl:621 [inlined]
  [46] #remake#962
     @ ~/.julia/packages/SciMLBase/Zzjmy/src/remake.jl:618 [inlined]
  [47] remake
     @ ~/.julia/packages/SciMLBase/Zzjmy/src/remake.jl:617 [inlined]
  [48] loss
     @ repro.jl:58 [inlined]
  [49] loss
     @ repro.jl:0 [inlined]
  [50] augmented_julia_loss_60368_inner_13wrap
     @ repro.jl:0
  [51] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:8175 [inlined]
  [52] enzyme_call
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7588 [inlined]
  [53] AugmentedForwardThunk
     @ ~/.julia/packages/Enzyme/wlmJK/src/compiler.jl:7534 [inlined]
  [54] autodiff
     @ ~/.julia/packages/Enzyme/wlmJK/src/Enzyme.jl:513 [inlined]
  [55] autodiff
     @ ~/.julia/packages/Enzyme/wlmJK/src/Enzyme.jl:643 [inlined]
  [56] macro expansion
     @ ~/.julia/packages/Enzyme/wlmJK/src/sugar.jl:291 [inlined]
  [57] gradient
     @ ~/.julia/packages/Enzyme/wlmJK/src/sugar.jl:278 [inlined]
  [58] gradient
     @ ~/.julia/packages/DifferentiationInterface/seMaz/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl:237 [inlined]
  [59] gradient(f::typeof(loss), backend::AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, x::Vector{Float64}, contexts::Constant{Tuple{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{Vector{ModelingToolkitBase.DiffCacheAllocatorAPIWrapper{Float64}}}}, ODEFunction{true, SciMLBase.AutoDespecialize, ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:___mtkunknowns___, :___mtkparameters___, :__argₛᵧₘ8912177291400142833), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0x8b376b31, 0x5e9e5a0d, 0xe002801d, 0xcb558804, 0x62e6ffb2), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__argₛᵧₘ14743185283305486891, :___mtkunknowns___, :___mtkparameters___, :__argₛᵧₘ8912177291400142833), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0x262a6e1c, 0xd91a41ee, 0xc3251555, 0x9af809af, 0x7b208f3b), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, ModelingToolkitBase.ObservedFunctionCache{System, Nothing}, Nothing, System, Union{Nothing, SciMLBase.OverrideInitData}, Union{Nothing, SciMLBase.ODENLStepData}}, @Kwargs{}, SciMLBase.StandardODEProblem}, SymbolicIndexingInterface.FullSetter{Nothing, SymbolicIndexingInterface.OOPSetter{false, System, Vector{ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, Int64}}}, Vector{Int64}, Vector{Int64}}, Tsit5{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), FastBroadcast.Serial}}})
     @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/seMaz/src/first_order/gradient.jl:63
  [60] top-level scope
     @ repro.jl:72
  [61] include(mod::Module, _path::String)
     @ Base ./Base.jl:309
  [62] exec_options(opts::Base.JLOptions)
     @ Base ./client.jl:344
  [63] _start()
     @ Base ./client.jl:577
in expression starting at repro.jl:72

Reproducer

# Run (the three cases below):
#   julia repro.jl                                               # inline, Tsit5
#   REPRO_SOLVER=Rodas5P julia repro.jl                          # inline, Rodas5P
#   REPRO_INLINE=false REPRO_SOLVER=Rodas5P julia repro.jl       # no inline solve
# Without the inline solve the model is a mass-matrix DAE, so it needs Rodas5P.

using ModelingToolkit, OrdinaryDiffEqTsit5, OrdinaryDiffEqRosenbrock, SciMLSensitivity,
      SciMLBase
using OrdinaryDiffEqNonlinearSolve # for the initialization
using ModelingToolkit: t_nounits as t, D_nounits as D
import ModelingToolkitTearing as MTKTearing
using SymbolicIndexingInterface: setsym_oop
using Enzyme, DifferentiationInterface, ADTypes, ForwardDiff

inline = get(ENV, "REPRO_INLINE", "true") == "true"
solver = get(ENV, "REPRO_SOLVER", "Tsit5") == "Tsit5" ? Tsit5() :
         Rodas5P(autodiff = AutoFiniteDiff())
println("inline = $inline, solver = $(nameof(typeof(solver)))")

# (y, z, w) form a linear system in which p enters the matrix; it is diagonally
# dominant for t ∈ [0, 1], so it stays nonsingular
@parameters p=1.0 k=2.0
@variables x(t)=1.0 y(t) [guess = 1.0] z(t) [guess = 1.0] w(t) [guess = 1.0] q(t)=1.0
eqs = [D(x) ~ 2t + 1,
    (5 + p * t) * y + x * z + w ~ 4,
    4y + 9z + 2w ~ 7,
    2x * y + 3t * z + 12w ~ 10,
    D(q) ~ k * w]
ralg = MTKTearing.DefaultReassembleAlgorithm(; inline_linear_sccs = inline)
sys = mtkcompile(System(eqs, t; name = :scc); reassemble_alg = ralg)
println("inline linear systems: ", length(MTKTearing.inline_linear_systems(sys)))

prob = ODEProblem(sys, [], (0.0, 1.0))
println("initialization problem: ",
    nameof(typeof(prob.f.initialization_data.initializeprob)))

function loss(x, ps)
    prob0, set_x, alg = ps
    u0, p = set_x(prob0, x)
    sol = solve(remake(prob0; u0, p), alg; saveat = 0.1, abstol = 1e-10,
        reltol = 1e-10)
    sum(abs2, reduce(vcat, sol.u))
end
ps = (prob, setsym_oop(prob, [p, k]), solver)
x0 = [0.9, 1.6]
let (u0, p) = ps[2](prob, x0)
    sol = solve(remake(prob; u0, p), solver; saveat = 0.1, abstol = 1e-10, reltol = 1e-10)
    println("primal retcode: ", sol.retcode, ", t_end = ", sol.t[end])
end

ref = DifferentiationInterface.gradient(loss, AutoForwardDiff(), x0, Constant(ps))
println("ForwardDiff: ", ref)
backend = AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse))
g = DifferentiationInterface.gradient(loss, backend, x0, Constant(ps))
println("Enzyme:      ", g)
println("isapprox: ", isapprox(g, ref; rtol = 1e-6))

(The script that produced the table also printed the package versions listed below; that loop is left out here.)

case inline systems primal ForwardDiff Enzyme
inline, Tsit5() 1 Success, t = 1.0 [0.296018, 12.765222] error above
inline, Rodas5P(autodiff = AutoFiniteDiff()) 1 Success, t = 1.0 [0.296018, 12.765222] error above
no inline solve, Rodas5P(autodiff = AutoFiniteDiff()) 0 Success, t = 1.0 [0.0281376, 12.765222] [0.0281376, 12.765222], isapprox: true

All three use an SCCNonlinearProblem initialization. The two inline cases fail at the same place, inside remake, so the solver plays no part in the failure.

The third case is a weaker control than it looks. Without the inline solve, the model is a mass-matrix DAE whose unknowns include y, z, w, so the loss (which sums sol.u) is a different function; that is why the ∂/∂p entries differ. It shows that Enzyme handles the same equations when they are not compiled to the inline solve. It does not check an inline-solve gradient.

Versions

Julia 1.13.3; ModelingToolkitTearing 1.20.9, ModelingToolkitBase 1.79.1, ModelingToolkit 11.46.1, LinearSolve 5.13.1, Enzyme 0.13.218 (Enzyme_jll 0.0.305), SciMLBase 3.58.0, SciMLSensitivity 7.119.14, OrdinaryDiffEqCore 4.19.1, OrdinaryDiffEqTsit5 2.1.5, OrdinaryDiffEqRosenbrock 2.7.5, OrdinaryDiffEqNonlinearSolve 2.9.3, NonlinearSolve 4.32.0, DifferentiationInterface 0.7.21, SymbolicIndexingInterface 0.3.56, ADTypes 1.24.0, ForwardDiff 1.4.7.

Possible directions (not tried)

  • Mark the lookup inactive (EnzymeRules.inactive on get_inline_linsolve_cache). The cache would then reach LinearSolve's solve! rule as a constant, and that rule doesn't support one: its augmented primal reads linsolve.dval.A, and its reverse pass asserts !(linsolve isa Const) (LinearSolve 5.13.1, ext/LinearSolveEnzymeExt.jl). So this alone probably just moves the failure.
  • An EnzymeRules rule for numeric_ldiv!(A, b) keeps the cache out of what Enzyme differentiates. Two things such a rule has to handle:
    • b is overwritten with the solution u and returned. The shadow of b holds ū on entry to the reverse pass. The rule has to read ū from it, then write b̄ = λ back into the same shadow (where λ = A⁻ᵀ ū), and accumulate Ā -= λ uᵀ. The output seed must not be counted twice.
    • The cache is shared by all solve sites of the same type and size, and A, b are scratch buffers rewritten on every call. So the reverse pass cannot rely on what the cache holds by then. The augmented primal has to save this call's u and A, or its factorization.

Written by Claude Code (claude-opus-5-5) on behalf of Sebastian. The reproducer and logs were run as shown; an independent model review (codex) was done on the draft.

Lingua principale
Julia
Stelle
8
Fork
14
Merge medio
2g 14h
PR unite (30g)
4

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di JuliaComputing/StateSelection.jl

Tutte le issue di JuliaComputing/StateSelection.jl

Issue simili

Altre issue su Julia

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.