Recently, I have been using the checkpointing method to reduce the memory required for automatic differentiation.
Therefore, we use Enzyme.jl and Checkpointing.jl.
However, when I use the latest version of Enzyme.jl, v0.12.33, an error occurs.
The version of Checkpointing.jl is v0.9.4.
The calculations work fine until v0.12.27 of Enzyme.jl.
What change in Enzyme.jl is causing this?
Below is a sample program and execution results of Checkpointing.jl.
# Explicit 1D heat equation
using Checkpointing
using Enzyme
using Plots
mutable struct Heat
Tnext::Vector{Float64}
Tlast::Vector{Float64}
n::Int
λ::Float64
tsteps::Int
step::Int64
end
function advance(heat)
next = heat.Tnext
last = heat.Tlast
λ = heat.λ
n = heat.n
for i in 2:(n-1)
next[i] = last[i] + λ*(last[i-1]-2*last[i]+last[i+1])
end
return nothing
end
function sumheat(heat::Heat, chkpscheme::Scheme, tsteps::Int64)
# AD: Create shadow copy for derivatives
@checkpoint_struct chkpscheme heat for i in 1:tsteps
heat.Tlast .= heat.Tnext
advance(heat)
end
return reduce(+, heat.Tnext)
end
function heat(scheme::Scheme, tsteps::Int)
n = 100
Δx=0.1
Δt=0.001
# Select μ such that λ ≤ 0.5 for stability with μ = (λ*Δt)/Δx^2
λ = 0.5
# Create object from struct. tsteps is not needed for a for-loop
heat = Heat(zeros(n), zeros(n), n, λ, tsteps,1)
# Shadow copy for Enzyme
dheat = Heat(zeros(n), zeros(n), n, λ, tsteps,1)
# Boundary conditions
heat.Tnext[1] = 20.0
heat.Tnext[end] = 0
# Compute gradient
autodiff(Enzyme.ReverseWithPrimal, sumheat, Duplicated(heat, dheat), Const(scheme), Const(tsteps))
return heat.Tnext, dheat.Tnext[2:end-1]
end
tsteps = 5
T, dT = heat(Revolve{Heat}(tsteps,4; verbose=1), tsteps)
# Plot function values
plot(T)
# Plot gradient with respect of sum(T[end]) with respect to T[1].
plot(dT)
[ Info: Revolve: Number of checkpoints: 4
[ Info: Revolve: Number of steps: 5
[ Info: Prediction:
[ Info: Forward steps : 4
[ Info: Overhead factor : 0.8
[ Info: Size per checkpoint: 1.688 KiB
[ Info: Revolve: First Uturn
[ Info: Size of total storage: 6.828 KiB
ERROR: Function argument passed to autodiff cannot be proven readonly.
If the the function argument cannot contain derivative data, instead call autodiff(Mode, Const(f), ...)
See https://enzyme.mit.edu/index.fcgi/julia/stable/faq/#Activity-of-temporary-storage for more information.
The potentially writing call is %67 = call nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* noundef nonnull @ijl_apply_generic, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4829008240 to {}*) to {} addrspace(10)*), {} addrspace(10)* nonnull %52, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 4367827296 to {}*) to {} addrspace(10)*)) #22, !dbg !202, using %52 = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %51 unordered, align 8, !dbg !202, !tbaa !34, !alias.scope !38, !noalias !41
Stacktrace:
[1] #16
@ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:182
[2] #16
@ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:0
Stacktrace:
[1] #16
@ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:182 [inlined]
[2] diffejulia__16_7917_inner_1wrap
@ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:0
[3] macro expansion
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:7099 [inlined]
[4] enzyme_call
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6708 [inlined]
[5] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6585 [inlined]
[6] autodiff
@ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:320 [inlined]
[7] autodiff
@ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:348 [inlined]
[8] autodiff
@ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:329 [inlined]
[9] rev_checkpoint_struct_for
@ ~/.julia/packages/Checkpointing/CEFSa/src/Schemes/Revolve.jl:471
[10] macro expansion
@ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:176 [inlined]
[11] sumheat
@ ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:31 [inlined]
[12] diffejulia_sumheat_7819wrap
@ ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:0
[13] macro expansion
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:7099 [inlined]
[14] enzyme_call
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6708 [inlined]
[15] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6585 [inlined]
[16] autodiff
@ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:320 [inlined]
[17] autodiff
@ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:348 [inlined]
[18] autodiff
@ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:329 [inlined]
[19] heat(scheme::Revolve{Heat}, tsteps::Int64)
@ Main ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:55
[20] top-level scope
@ ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:60