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

オープン
  • 何か問題がありますか？
#183 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
48/100
issue の種類
バグ
明瞭さ
おおむね明確
活発さ
活発
技術スタック
julia
領域
backend-api-design

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

説明

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.

主要言語
Julia
スター
8
フォーク
14
平均マージ
2日 14時間
マージ済み PR（30日）
4

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

JuliaComputing/StateSelection.jl のほかの issue

JuliaComputing/StateSelection.jl の issue をすべて見る

似ている issue

Julia の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。