Add integration tests for Checkpointing (#2811)

Author: vchuravy

.github/workflows/Integration.yml

@@ -46,6 +46,7 @@ jobs:
           - linux-x86-n2-32
         package:
           - Bijectors
+          - Checkpointing
           - DifferentiationInterface
           - Distributions
           - DynamicExpressions

test/integration/Checkpointing/Project.toml

@@ -0,0 +1,14 @@
+[deps]
+Checkpointing = "eb46d486-4f9c-4c3d-b445-a617f2a2f1ca"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[compat]
+Checkpointing = "0.11"
+
+[sources.Enzyme]
+path = "../../.."
+
+[sources.EnzymeCore]
+path = "../../../lib/EnzymeCore"

test/integration/Checkpointing/runtests.jl

@@ -0,0 +1,59 @@
+using Test
+using Enzyme
+using Checkpointing
+using LinearAlgebra
+
+# Explicit 1D heat equation
+mutable struct Heat
+    Tnext::Vector{Float64}
+    Tlast::Vector{Float64}
+    n::Int
+    λ::Float64
+    tsteps::Int
+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_checkpoint chkpscheme 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)
+    # Shadow copy for Enzyme
+    dheat = Heat(zeros(n), zeros(n), n, λ, tsteps)
+
+    # 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
+
+T, dT = heat(Revolve(4), 500)
+@test norm(T) == 66.21987468492061
+@test norm(dT) == 6.970279349365908