ITADN

Nested reverse-mode (set_runtime_activity) over a const×active BLAS gemv corrupts/zeros/segfaults the outer gradient ("freeing without malloc")

#3139OpenChrisRackauckas-Claude 创建于 2026-06-03
## Summary When an outer `Enzyme.autodiff(set_runtime_activity(Reverse), …)` differentiates a function whose body runs a **nested** `Enzyme.gradient(set_runtime_activity(Reverse), …)` over a dense BLAS `gemv` (constant matrix × active vector), Enzyme mismanages the nested BLAS adjoint tape (codegen prints `freeing without malloc …`) and the **outer gradient is silently wrong** — exactly zero, 2×-off, or a segfault depending on how the constant matrix is held. Single-level reverse of the same inner function is correct; the corruption appears only under nesting, and only when the nested reverse pass goes through dense BLAS/LAPACK. Versions: **Enzyme 0.13.152**, Julia **1.11.9**, OpenBLAS (ILP64). ## MWE 1 — silently returns zero ```julia using Enzyme, LinearAlgebra const M = [1.0 0.2 0.0; 0.0 1.0 0.1; 0.3 0.0 1.0] inner(t) = sum((M * t) .^ 2) # const M × active t -> BLAS gemv g(p) = sum(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(inner), p)[1]) # nested reverse dp = zero([1.0, 2.0, 3.0]) Enzyme.autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(g), Enzyme.Active, Enzyme.Duplicated([1.0, 2.0, 3.0], dp)) @show dp # [0.0, 0.0, 0.0] @show 2 .* (transpose(M)*M * ones(3)) # correct: [3.18, 2.68, 2.82] ``` Output: ``` freeing without malloc %24 = extractvalue { double*, ... } %tapeArg, 2 ... dp = [0.0, 0.0, 0.0] ``` ## MWE 2 — segfault (constant matrix held in a struct, captured by the inner closure) ```julia using Enzyme, LinearAlgebra struct Cfg; M::Matrix{Float64}; end make_inner(cfg::Cfg) = t -> sum((cfg.M * t) .^ 2) function g(p, cfg) ig = Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(make_inner(cfg)), p)[1] return sum(p .^ 2) + sum(ig .* p) end dp = zero([1.0,2.0,3.0]) Enzyme.autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(g), Enzyme.Active, Enzyme.Duplicated([1.0,2.0,3.0], dp), Enzyme.Const(Cfg([1.0 0.2 0.0; 0.0 1.0 0.1; 0.3 0.0 1.0]))) ``` → `signal 11 (Segmentation fault)`, preceded by the same `freeing without malloc` warnings (deterministic, also crashes inside `_generic_matvecmul!`). ## Variants (all reproduced, same `freeing without malloc` signature) - **2×-wrong:** `g(p) = sum(Enzyme.gradient(Reverse, Const(inner), p)[1] .* p)` with the const-global `M` returns `[4.78, 5.16, 7.06]` instead of `[9.56, 10.32, 14.12]` (exactly half). - **Correct when not nested:** a single-level `Enzyme.gradient(Reverse, Const(inner), p)` is correct; the bug needs the outer reverse pass. - **Non-BLAS inner functions don't corrupt:** smooth scalar inners, an explicit (non-BLAS) 2×2 solve, or an elementwise const-vector inner all give correct nested gradients. The necessary ingredient is dense BLAS in the nested reverse pass mixing constant + active memory. - **Related, possibly separate:** if the nested reverse path uses a LAPACK `lu`/`\` instead of `gemv`, it fails at compile time with `UndefVarError: source_sig not defined in Enzyme.Compiler` (from `Enzyme/src/typeutils/jltypes.jl:206` — `source_sig` isn't a parameter of the enclosing function, so a real `classify_arguments` source-pointer-type mismatch is masked by the UndefVarError). ## Context This surfaced through SciML/SciMLSensitivity.jl#1469: differentiating an MTK DAE-initialization adjoint with Enzyme on the outside computes a nested init-sensitivity gradient that routes through NonlinearSolve factorizations / the semi-explicit DAE adjoint's linear solves; the corrupted/zeroed nested tape produces a downstream `LinearAlgebra.SingularException`. The SciML-side workaround was to avoid the nesting. The MWEs above reduce it to Enzyme + LinearAlgebra only. Happy to provide the additional variant scripts.
3 条评论