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

Đang mở
  • Có gì không ổn?
#183 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
48/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
julia
Lĩnh vực
backend-api-design

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Julia
Star
8
Fork
14
Merge trung bình
2 ngày 14 giờ
Pull request đã merge (30 ngày)
4

Chuẩn bị môi trường

Dự án này không cung cấp dev container, Dockerfile hay hướng dẫn đóng góp, nên bạn cần tự thiết lập môi trường: hãy bắt đầu từ README và xem hướng dẫn đóng góp lần đầu của chúng tôi để biết các bước chung.

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của JuliaComputing/StateSelection.jl

Tất cả issue của JuliaComputing/StateSelection.jl

Issue tương tự

Thêm issue về Julia

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.