EnzymeAD / Enzyme.jl

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

Runtime activity required for invoke #407

Closed ForceBru closed 2 weeks ago

ForceBru commented 2 years ago

Code

import Pkg; Pkg.status()

import Random
import Enzyme

const AV = AbstractVector{T} where T

# ===== Set up objective function =====
normal_pdf(x::Real, mean::Real, var::Real) =
    exp(-(x - mean)^2 / (2var)) / sqrt(2π * var)

function mixture_loglikelihood(params::AV{<:Real}, data::AV{<:Real})::Real
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[(K+1):2K], params[(2K+1):end]

    sum(
        sum(
            weight * normal_pdf(x, mean, std)
            for (weight, mean, std) in zip(weights, means, stds)
        ) |> log
        for x in data
    )
end

rnd = Random.MersenneTwister(42)
data = randn(rnd, 500)
params0 = [[0.2, 0.8]; [-1, 1]; [.3, .6]]
objective = params -> mixture_loglikelihood(params, data)

# ===== Compute gradient =====
# 1. From docs (https://docs.juliahub.com/Enzyme/G1p5n/0.10.4/api/#Enzyme.gradient!-Tuple{Enzyme.ReverseMode,%20Any,%20Any,%20Any})
let
    f(x) = x[1]*x[2]

    dx = [0.0, 0.0]
    Enzyme.gradient!(Enzyme.Reverse, dx, f, [2.0, 3.0])
    @show dx
end

# 2. My function
let
    grad_storage = similar(params0)
    Enzyme.gradient!(Enzyme.Reverse, grad_storage, objective, params0)
    @show grad_storage
end

Output

Status `~/test/autodiff_bench/Project.toml`
  [6e4b80f9] BenchmarkTools v1.3.1
  [7da242da] Enzyme v0.10.4
  [f6369f11] ForwardDiff v0.10.32
  [37e2e3b7] ReverseDiff v1.14.1
  [e88e6eb3] Zygote v0.6.43
dx = [3.0, 2.0]
ERROR: LoadError: UndefRefError: access to undefined reference
Stacktrace:
 [1] unsafe_convert
   @ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4277 [inlined]
 [2] macro expansion
   @ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4520 [inlined]
 [3] enzyme_call(::Ptr{Nothing}, ::Type{Enzyme.Compiler.AugmentedForwardThunk}, ::Type{Val{0x0000000000000001}}, ::Val{true}, ::Type{Tuple{Enzyme.Duplicated{Vector{Float64}}, Enzyme.Const{Vector{Float64}}}}, ::Type{Enzyme.Const{Real}}, ::typeof(mixture_loglikelihood), ::Nothing, ::Enzyme.Duplicated{Vector{Float64}}, ::Enzyme.Const{Vector{Float64}})
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4316
 [4] (::Enzyme.Compiler.AugmentedForwardThunk{typeof(mixture_loglikelihood), Enzyme.Const{Real}, Tuple{Enzyme.Duplicated{Vector{Float64}}, Enzyme.Const{Vector{Float64}}}, Val{0x0000000000000001}, Nothing, Val{true}()})(::Enzyme.Duplicated{Vector{Float64}}, ::Vararg{Any})
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4307
in expression starting at /Users/forcebru/test/autodiff_bench/Enzyme_bug.jl:39

Other autodiff packages (I tried ForwardDiff, ReverseDiff and Zygote) don't have any issues differentiating objective.

Versions

julia> versioninfo()
Julia Version 1.8.0-rc3
Commit 33f19bcbd25 (2022-07-13 19:10 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin21.4.0)
  CPU: 4 × Intel(R) Core(TM) i5-3330S CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, ivybridge)
  Threads: 1 on 4 virtual cores
ForceBru commented 2 years ago

I changed mixture_loglikelihood to use index access and remove @views, like in f(x) from the docs:

function mixture_loglikelihood(params::AV{<:Real}, data::AV{<:Real})::Real
    K = length(params) ÷ 3
    weights, means, stds = params[1:K], params[(K+1):2K], params[(2K+1):end]

    sum(
        sum(
            weights[k] * normal_pdf(data[n], means[k], stds[k])
            for k in 1:K
        ) |> log
        for n in eachindex(data)
    )
end

However, I'm still getting the same error:

warning: didn't implement memmove, using memcpy as fallback which can result in errors
warning: didn't implement memmove, using memcpy as fallback which can result in errors
ERROR: LoadError: UndefRefError: access to undefined reference
Stacktrace:
 [1] unsafe_convert
   @ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4277 [inlined]
 [2] macro expansion
   @ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4520 [inlined]
 [3] enzyme_call(::Ptr{Nothing}, ::Type{Enzyme.Compiler.AugmentedForwardThunk}, ::Type{Val{0x0000000000000001}}, ::Val{true}, ::Type{Tuple{Enzyme.Duplicated{Vector{Float64}}, Enzyme.Const{Vector{Float64}}}}, ::Type{Enzyme.Const{Real}}, ::typeof(mixture_loglikelihood), ::Nothing, ::Enzyme.Duplicated{Vector{Float64}}, ::Enzyme.Const{Vector{Float64}})
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4316
 [4] (::Enzyme.Compiler.AugmentedForwardThunk{typeof(mixture_loglikelihood), Enzyme.Const{Real}, Tuple{Enzyme.Duplicated{Vector{Float64}}, Enzyme.Const{Vector{Float64}}}, Val{0x0000000000000001}, Nothing, Val{true}()})(::Enzyme.Duplicated{Vector{Float64}}, ::Vararg{Any})
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4307
in expression starting at /Users/forcebru/test/autodiff_bench/Enzyme_bug.jl:42
vchuravy commented 2 years ago

@wsmoses was your comment in https://github.com/EnzymeAD/Enzyme.jl/issues/347#issuecomment-1212435054 w.r.t this issue?

vchuravy commented 2 years ago

On https://github.com/EnzymeAD/Enzyme.jl/pull/408 I see:

vchuravy@odin ~/s/Enzyme (vc/anonymous)> julia --project test.jl
dx = [3.0, 2.0]
ERROR: LoadError: TypeError: 
signal (11): Segmentation fault
in expression starting at none:0
ijl_subtype_env at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/subtype.c:1857
ijl_isa at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/subtype.c:2114
showerror at ./errorshow.jl:69
#showerror#861 at ./errorshow.jl:88
showerror##kw at ./errorshow.jl:86
unknown function (ip: 0x7f90e850c242)
_jl_invoke at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2358 [inlined]
ijl_apply_generic at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2540
#showerror#862 at ./errorshow.jl:96
showerror##kw at ./errorshow.jl:94
unknown function (ip: 0x7f90e850ab12)
_jl_invoke at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2358 [inlined]
ijl_apply_generic at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2540
show_exception_stack at ./errorshow.jl:878
display_error at ./client.jl:103
unknown function (ip: 0x7f90e850a401)
_jl_invoke at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2358 [inlined]
ijl_apply_generic at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2540
display_error at ./client.jl:106
jfptr_display_error_40695.clone_1 at /home/vchuravy/.julia/juliaup/julia-1.8.0-rc3+0.x64/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2358 [inlined]
ijl_apply_generic at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2540
jl_apply at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/julia.h:1838 [inlined]
jl_f__call_latest at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/builtins.c:774
#invokelatest#2 at ./essentials.jl:729 [inlined]
invokelatest at ./essentials.jl:726 [inlined]
exec_options at ./client.jl:306
_start at ./client.jl:522
jfptr__start_45041.clone_1 at /home/vchuravy/.julia/juliaup/julia-1.8.0-rc3+0.x64/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2358 [inlined]
ijl_apply_generic at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/gf.c:2540
jl_apply at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/julia.h:1838 [inlined]
true_main at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/jlapi.c:575
jl_repl_entrypoint at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/jlapi.c:719
main at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/cli/loader_exe.c:59
unknown function (ip: 0x7f9187b0a2cf)
__libc_start_main at /usr/lib/libc.so.6 (unknown line)
unknown function (ip: 0x401098)
Allocations: 79647333 (Pool: 79544916; Big: 102417); GC: 93
vchuravy commented 2 years ago

Ok some bread-crumbs

Thread 1 received signal SIGSEGV, Segmentation fault.
0x00007f8a22c977d0 in jl_is_type (v=0x0) at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/julia.h:1255
1255    /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/julia.h: No such file or directory.
(rr) up
#1  jl_f_issubtype (F=<optimized out>, args=0x7ffe851c1df0, nargs=<optimized out>)
    at /cache/build/default-amdci5-5/julialang/julia-release-1-dot-8/src/builtins.c:507
(rr) p jl_(args[0])
#<null>
$3 = void
(rr) p jl_(args[1])
AbstractFloat
#2  0x00007f8a22c3d404 in runtime_invoke_rev () at /home/vchuravy/src/Enzyme/src/compiler.jl:630
630                 if typeof(p) <: AbstractFloat || typeof(p) <: Complex{<:AbstractFloat}

So looks like typeof(p) is corrupted.

vchuravy commented 2 years ago

@wsmoses Can you take a look?

vchuravy@odin ~/s/Enzyme (vc/anonymous)> ~/builds/julia-dbg/julia --project=. test.jl
dx = [3.0, 2.0]
julia: /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/ADT/DenseMap.h:1250: bool llvm::operator==(const llvm::DenseMapIterator<llvm::ValueMapCallbackVH<const llvm::Value*, InvertedPointerVH, llvm::ValueMapConfig<const llvm::Value*, llvm::sys::SmartMutex<false> > >, InvertedPointerVH, llvm::DenseMapInfo<llvm::ValueMapCallbackVH<const llvm::Value*, InvertedPointerVH, llvm::ValueMapConfig<const llvm::Value*, llvm::sys::SmartMutex<false> > > >, llvm::detail::DenseMapPair<llvm::ValueMapCallbackVH<const llvm::Value*, InvertedPointerVH, llvm::ValueMapConfig<const llvm::Value*, llvm::sys::SmartMutex<false> > >, InvertedPointerVH>, false>&, const llvm::DenseMapIterator<llvm::ValueMapCallbackVH<const llvm::Value*, InvertedPointerVH, llvm::ValueMapConfig<const llvm::Value*, llvm::sys::SmartMutex<false> > >, InvertedPointerVH, llvm::DenseMapInfo<llvm::ValueMapCallbackVH<const llvm::Value*, InvertedPointerVH, llvm::ValueMapConfig<const llvm::Value*, llvm::sys::SmartMutex<false> > > >, llvm::detail::DenseMapPair<llvm::ValueMapCallbackVH<const llvm::Value*, InvertedPointerVH, llvm::ValueMapConfig<const llvm::Value*, llvm::sys::SmartMutex<false> > >, InvertedPointerVH>, false>&): Assertion `(!LHS.Ptr || LHS.isHandleInSync()) && "handle not in sync!"' failed.

signal (6): Aborted
in expression starting at /home/vchuravy/src/Enzyme/test.jl:40
unknown function (ip: 0x7f98a8a204dc)
gsignal at /usr/lib/libc.so.6 (unknown line)
abort at /usr/lib/libc.so.6 (unknown line)
unknown function (ip: 0x7f98a89ba45b)
__assert_fail at /usr/lib/libc.so.6 (unknown line)
operator== at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/ADT/DenseMap.h:1250
operator!= at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/ADT/DenseMap.h:1259 [inlined]
operator!= at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/ValueMap.h:368 [inlined]
visitCallInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:8526
delegateCallInst at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/InstVisitor.h:302 [inlined]
visitCall at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/Instruction.def:209 [inlined]
visit at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/Instruction.def:209
visit at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/InstVisitor.h:112 [inlined]
CreateAugmentedPrimal at /workspace/srcdir/Enzyme/enzyme/Enzyme/EnzymeLogic.cpp:2201
EnzymeCreateAugmentedPrimal at /workspace/srcdir/Enzyme/enzyme/Enzyme/CApi.cpp:474
EnzymeCreateAugmentedPrimal at /home/vchuravy/src/Enzyme/src/api.jl:147

With:

vchuravy@odin ~/s/Enzyme (vc/anonymous) [SIGABRT]> cat ~/builds/julia-dbg/Make.user
FORCE_ASSERTIONS=1
LLVM_ASSERTIONS=1
override JULIA_BUILD_MODE=debug
wsmoses commented 1 year ago

@vchuravy the not in sync error was fixed, what happens now

wsmoses commented 1 year ago
import Enzyme
Enzyme.API.printall!(true)

@noinline function myfoldl_impl(itr, cond)
    # Unroll the while loop once; if init is known, the call to op may
    # be evaluated at compile time
    y = iterate(itr)
    if cond
        return Base._InitialValue()
    else
        return itr[1]
    end
end

function mixture_loglikelihood(data::Vector{Float64})::Real
    myfoldl_impl(data, false)
end

data = ones(Float64, 500)
ddata = zeros(Float64, 500)

forward, pullback = Enzyme.Compiler.thunk(mixture_loglikelihood, nothing, Enzyme.Active, Tuple{Enzyme.Duplicated{Vector{Float64}}}, Val(Enzyme.API.DEM_ReverseModeGradient), Val(1))
dup = Enzyme.Duplicated(data, ddata)
@show forward(dup)
wsmoses commented 1 year ago

After fixing the bug in https://github.com/EnzymeAD/Enzyme.jl/pull/522 this is now a pure segmentation fault for the original code above.

wsmoses commented 1 year ago

New segmentation fault:

wmoses@beast:~/git/Enzyme.jl (genericninit) $ cat ur2.jl 
import Enzyme
Enzyme.API.printall!(true)

function mymapfoldl_impl(f::F, op::OP, nt, itr) where {F,OP}
    op′, itr′ = Base._xfadjoint(Base.BottomRF(op), Base.Generator(f, itr))
    return my_foldl_impl(op′, nt, itr′)
end

function my_foldl_impl(op::OP, init, itr) where {OP}
    # Unroll the while loop once; if init is known, the call to op may
    # be evaluated at compile time
    y = iterate(itr)
    y === nothing && return init
    return op(init, y[1])
end

function mixture_loglikelihood(weights::Vector{Float64}, data::Vector{Float64})::Float64
    itr1 = (weight for (weight, mean) in zip(weights, weights))

    itr2 = (mymapfoldl_impl(Base.identity, Base.add_sum, Base._InitialValue(), itr1) for x in data)

    mymapfoldl_impl(Base.identity, Base.add_sum, Base._InitialValue(), itr2)
end

data = ones(Float64, 1)

weights = [0.2]
dweights = [0.0]

@show mixture_loglikelihood(weights, data)

Enzyme.autodiff(Enzyme.Reverse, mixture_loglikelihood, Enzyme.Duplicated(weights, dweights), Enzyme.Const(data))
@show dweights
wsmoses commented 1 year ago

Now functions without segfaulting, but requires runtimeActivity for correctness:

import Enzyme

Enzyme.API.runtimeActivity!(true)

function mixture_loglikelihood(weights::Vector{Float64}, data::Vector{Float64})::Float64
    sum(
        sum(
            x * weight
            for (weight, mean, std) in zip(weights, weights, weights)
        ) 
        for x in data
    )
end

data = [-0.6, -0.5]
params0 = [0.2, 0.8]
objective = params -> mixture_loglikelihood(params, data)

let
    # grad_storage = zero(params0)
    # Enzyme.gradient!(Enzyme.Reverse, grad_storage, objective, params0)
    # @show grad_storage

@show data
    grad_storage = zero(params0)
    Enzyme.autodiff(Enzyme.Reverse, objective, Enzyme.Active, Enzyme.Duplicated(params0, grad_storage))
    @show grad_storage

@show data
    grad_storage = zero(params0)
    Enzyme.autodiff(Enzyme.Reverse, mixture_loglikelihood, Enzyme.Active, Enzyme.Duplicated(params0, grad_storage), Enzyme.Const(data))
    @show grad_storage

@show data
end
# without
wmoses@beast:~/git/Enzyme.jl (genericninit) $ ./julia-1.8.2/bin/julia --project ur3.jl
Project Enzyme v0.10.12
Status `~/git/Enzyme.jl/Project.toml`
  [fa961155] CEnum v0.4.2
  [f151be2c] EnzymeCore v0.1.0
  [61eb1bfa] GPUCompiler v0.16.4
  [929cbde3] LLVM v4.14.0
  [d8793406] ObjectFile v0.3.7
  [7cc45869] Enzyme_jll v0.0.43+0
  [8f399da3] Libdl
  [37e2e46d] LinearAlgebra
  [de0858da] Printf
  [9a3f8284] Random
data = [-0.6, -0.5]
grad_storage = [-1.1, -1.1]
data = [0.4, 0.5]
grad_storage = [0.9, 0.9]
data = [1.4, 1.5]
wmoses@beast:~/git/Enzyme.jl (genericninit) $ ./julia-1.8.2/bin/julia --project ur3.jl
data = [-0.6, -0.5]
grad_storage = [-1.1, -1.1]
data = [-0.6, -0.5]
grad_storage = [-1.1, -1.1]
data = [-0.6, -0.5]
wsmoses commented 2 weeks ago

This is now expected behavior due to the closure. At one point in the future we may have fine-grained struct activity, but in any case closing.