Nested reverse-mode (set_runtime_activity) over a const×active BLAS gemv corrupts/zeros/segfaults the outer gradient ("freeing without malloc")
## 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 条评论