EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
455 stars 63 forks source link

Primal GC preserve end not postdominating begin -> Reverse-mode fails verification #877

Closed sethaxen closed 6 months ago

sethaxen commented 1 year ago

In this example, if the callable is mutated during its evaluation, then an error is raised during reverse-mode AD:

using Enzyme

struct MutatedCallable{T}
    x::T
end
function (c::MutatedCallable)(y)
    s = c.x'y
    c.x ./= s
    return s
end

c = MutatedCallable(randn(3))
∂c = MutatedCallable(zeros(3))
y = randn(3)
∂y = zeros(3)

autodiff(Forward, Duplicated(c, ∂c), Duplicated, Duplicated(y, ∂y))  # fine
autodiff(Reverse, Duplicated(c, ∂c), Active, Duplicated(y, ∂y))  # errors, see attached

dump.txt

vchuravy commented 1 year ago

In particular:

Instruction does not dominate all uses!
  %168 = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %.fca.0.extract, {} addrspace(10)* %".fca.0.extract'ipev", {} addrspace(10)* %1, {} addrspace(10)* %"'1")
  call void @llvm.julia.gc_preserve_end(token %168)
invertcblas_ddot64_.exit.i.thread:                ; preds = %invertjulia_MutatedCallable_3018_inner.exit
  %168 = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %.fca.0.extract, {} addrspace(10)* %".fca.0.extract'ipev", {} addrspace(10)* %1, {} addrspace(10)* %"'1")
  br label %invertL6.i

invertL6.i:                                       ; preds = %invertcblas_ddot64_.exit.i.thread, %invert.preheader.preheader
  call void @llvm.julia.gc_preserve_end(token %168)
  br label %invertentry

So it looks like we didn't split the edge from %invertcblas_ddot64_.exit.i.thread to invertL6.i: and have a second predecessor

wsmoses commented 9 months ago

Is fixed by https://github.com/EnzymeAD/Enzyme.jl/pull/1282