Make tape recording thread-safe (fixes #218)

.github/workflows/ci.yml

@@ -48,6 +48,8 @@ jobs:
       - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        env:
+          JULIA_NUM_THREADS: 4
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v7
         with:

src/api/Config.jl

@@ -105,7 +105,6 @@ function JacobianConfig(output::AbstractArray{D}, input::Tuple, tp::InstructionT
     return _JacobianConfig(cfg_input, cfg_output, tp)
 end
 
-# we dispatch on V<:Real here because InstructionTape is actually also an AbstractArray
 function JacobianConfig(output::AbstractArray{D}, input::AbstractArray{V}, tp::InstructionTape = InstructionTape()) where {D, V <: Real}
     cfg_input = track(similar(input), D, tp)
     cfg_output = track!(similar(output, TrackedReal{D, D, Nothing}), output, tp)

src/api/tape.jl

@@ -101,8 +101,8 @@ methods on each instruction in the tape.
 function CompiledTape(t::T) where {T <: AbstractTape}
     return CompiledTape{T}(
         t,
-        [FunctionWrapper{Nothing, Tuple{}}(ForwardExecutor(instruction)) for instruction in t.tape],
-        [FunctionWrapper{Nothing, Tuple{}}(ReverseExecutor(t.tape[i])) for i in length(t.tape):-1:1]
+        [FunctionWrapper{Nothing, Tuple{}}(ForwardExecutor(instruction)) for instruction in instructions(t.tape)],
+        [FunctionWrapper{Nothing, Tuple{}}(ReverseExecutor(instruction)) for instruction in Iterators.reverse(instructions(t.tape))]
     )
 end
 
@@ -204,6 +204,7 @@ See `ReverseDiff.gradient` for a description of acceptable types for `input`.
 function GradientTape(f, input, cfg::GradientConfig = GradientConfig(input))
     track!(cfg.input, input)
     tracked_ouput = f(cfg.input)
+    finish!(cfg.tape)
     return _GradientTape(f, cfg.input, tracked_ouput, cfg.tape)
 end
 
@@ -212,6 +213,7 @@ function GradientTape(f, input::Tuple, cfg::GradientConfig = GradientConfig(inpu
         track!(cfg.input[i], input[i])
     end
     tracked_output = f(cfg.input...)
+    finish!(cfg.tape)
     return _GradientTape(f, cfg.input, tracked_output, cfg.tape)
 end
 
@@ -234,6 +236,7 @@ See `ReverseDiff.jacobian` for a description of acceptable types for `input`.
 function JacobianTape(f, input, cfg::JacobianConfig = JacobianConfig(input))
     track!(cfg.input, input)
     tracked_ouput = f(cfg.input)
+    finish!(cfg.tape)
     return _JacobianTape(f, cfg.input, tracked_ouput, cfg.tape)
 end
 
@@ -242,6 +245,7 @@ function JacobianTape(f, input::Tuple, cfg::JacobianConfig = JacobianConfig(inpu
         track!(cfg.input[i], input[i])
     end
     tracked_output = f(cfg.input...)
+    finish!(cfg.tape)
     return _JacobianTape(f, cfg.input, tracked_output, cfg.tape)
 end
 
@@ -261,6 +265,7 @@ function JacobianTape(f!, output, input, cfg::JacobianConfig = JacobianConfig(ou
     track!(cfg.output, output, cfg.tape)
     track!(cfg.input, input)
     f!(cfg.output, cfg.input)
+    finish!(cfg.tape)
     return _JacobianTape(f!, cfg.input, cfg.output, cfg.tape)
 end
 
@@ -270,6 +275,7 @@ function JacobianTape(f!, output, input::Tuple, cfg::JacobianConfig = JacobianCo
         track!(cfg.input[i], input[i])
     end
     f!(cfg.output, cfg.input...)
+    finish!(cfg.tape)
     return _JacobianTape(f!, cfg.input, cfg.output, cfg.tape)
 end
 
@@ -296,5 +302,6 @@ function HessianTape(f, input, cfg::HessianConfig = HessianConfig(input))
     track!(ht.input, input)
     gt = GradientTape(f, ht.input, gcfg)
     seeded_reverse_pass!(ht.output, gt.output, gt.input, gt.tape)
+    finish!(jcfg.tape)
     return ht
 end

src/tape.jl

@@ -4,10 +4,75 @@
 
 abstract type AbstractInstruction end
 
-const InstructionTape = Vector{AbstractInstruction}
+mutable struct TapeNode
+    const instruction::AbstractInstruction
+    prev::Union{Nothing, TapeNode}
+end
+
+struct Finished end
+
+# A tape is first recorded, possibly from multiple threads at once, and then replayed.
+# Reading it while recording throws, and so does recording onto it after `finish!`;
+# `empty!` starts recording again.
+mutable struct InstructionTape
+    # while recording: the most recently recorded node (nodes link backwards);
+    # `Finished()` once recording has finished
+    @atomic last::Union{Nothing, TapeNode, Finished}
+    # recorded instructions in order, filled by `finish!`
+    const instructions::Vector{AbstractInstruction}
+end
+
+InstructionTape() = InstructionTape(nothing, AbstractInstruction[])
+
+# end recording; all tasks recording onto `tp` must have finished (e.g. via `@sync`)
+function finish!(tp::InstructionTape)
+    node = @atomicswap tp.last = Finished()
+    if node isa Finished
+        throw(ArgumentError("tape has already finished recording"))
+    end
+    while node !== nothing
+        push!(tp.instructions, node.instruction)
+        node = node.prev
+    end
+    reverse!(tp.instructions)
+    return tp
+end
+
+# the recorded instructions; throws unless recording has finished
+function instructions(tp::InstructionTape)
+    if !((@atomic tp.last) isa Finished)
+        throw(ArgumentError("tape is still recording; call `ReverseDiff.finish!` first"))
+    end
+    return tp.instructions
+end
+
+Base.length(tp::InstructionTape) = length(instructions(tp))
+
+function Base.empty!(tp::InstructionTape)
+    # `NULL_TAPE` never records, so it stays finished
+    if tp !== NULL_TAPE
+        empty!(tp.instructions)
+        @atomic tp.last = nothing
+    end
+    return tp
+end
+
+# undo the swap in `record!` so that the tape stays finished
+@noinline function throw_finished(tp::InstructionTape)
+    @atomic tp.last = Finished()
+    throw(ArgumentError("tape has finished recording; call `empty!` to record again"))
+end
 
-function record!(tp::InstructionTape, ::Type{InstructionType}, args...) where {InstructionType}
-    tp !== NULL_TAPE && push!(tp, InstructionType(args...))
+@inline function record!(tp::InstructionTape, ::Type{InstructionType}, args...) where {InstructionType <: AbstractInstruction}
+    if tp !== NULL_TAPE
+        node = TapeNode(InstructionType(args...), nothing)
+        # the swap orders concurrent recordings and returns the predecessor
+        prev = @atomicswap tp.last = node
+        if prev isa Finished
+            throw_finished(tp)
+        end
+        node.prev = prev
+    end
     return nothing
 end
 
@@ -75,7 +140,7 @@ end
 ##########
 
 function forward_pass!(tape::InstructionTape)
-    for instruction in tape
+    for instruction in instructions(tape)
         forward_exec!(instruction)
     end
     return nothing
@@ -85,8 +150,8 @@ end
 @noinline forward_exec!(instruction::SpecialInstruction) = special_forward_exec!(instruction)
 
 function reverse_pass!(tape::InstructionTape)
-    for i in length(tape):-1:1
-        reverse_exec!(tape[i])
+    for instruction in Iterators.reverse(instructions(tape))
+        reverse_exec!(instruction)
     end
     return nothing
 end
@@ -98,28 +163,41 @@ end
 # Pretty Printing #
 ###################
 
-# extra spaces here accomodates padding in show(::IO, ::AbstractInstruction)
-compactrepr(x::Tuple) = "(" * join(map(compactrepr, x), ",\n           ") * ")"
-compactrepr(x::AbstractArray) = length(x) < 5 ? match(r"\[.*?\]", repr(x)).match : summary(x)
-compactrepr(x) = repr(x)
-
-function Base.show(io::IO, instruction::AbstractInstruction, pad = "")
-    name = isa(instruction, ScalarInstruction) ? "ScalarInstruction" : "SpecialInstruction"
-    println(io, pad, "$(name)($(instruction.func)):")
-    println(io, pad, "  input:  ", compactrepr(instruction.input))
-    println(io, pad, "  output: ", compactrepr(instruction.output))
-    print(io, pad, "  cache:  ", compactrepr(instruction.cache))
+function Base.show(io::IO, instruction::AbstractInstruction)
+    print(io, nameof(typeof(instruction)), "(", instruction.func, ")")
+    return nothing
+end
+
+function Base.show(io::IO, ::MIME"text/plain", instruction::AbstractInstruction)
+    show(io, instruction)
+    ctx = IOContext(io, :compact => true, :limit => true)
+    print(ctx, ":\n  input:  ")
+    show(ctx, instruction.input)
+    print(ctx, "\n  output: ")
+    show(ctx, instruction.output)
+    print(ctx, "\n  cache:  ")
+    show(ctx, instruction.cache)
     return nothing
 end
 
 function Base.show(io::IO, tp::InstructionTape)
-    println(io, length(tp), "-element InstructionTape:")
-    i = 1
-    for instruction in tp
-        print(io, "$i => ")
-        show(io, instruction)
-        println(io)
-        i += 1
+    if (@atomic tp.last) isa Finished
+        print(io, length(tp), "-element InstructionTape")
+    else
+        print(io, "InstructionTape (recording)")
+    end
+    return nothing
+end
+
+# like arrays, so that long tapes are truncated
+function Base.show(io::IO, ::MIME"text/plain", tp::InstructionTape)
+    show(io, tp)
+    if (@atomic tp.last) isa Finished
+        instrs = instructions(tp)
+        if !isempty(instrs)
+            println(io, ":")
+            Base.print_array(io, instrs)
+        end
     end
-    return
+    return nothing
 end

src/tracked.jl

@@ -3,7 +3,7 @@
 #########
 
 const NULL_INDEX = typemin(Int)
-const NULL_TAPE = InstructionTape()
+const NULL_TAPE = finish!(InstructionTape())
 
 # TrackedReal #
 #-------------#

test/ChainRulesTests.jl

@@ -134,6 +134,7 @@ end
     xt = ReverseDiff.track(copy(x), tp)
     # record
     yt = f(xt)
+    ReverseDiff.finish!(tp)
     @test length(tp) == 3 # sum, broadcast+, broadcast*, but not `g`
 end
 

test/MacrosTests.jl

@@ -41,18 +41,18 @@ function test_forward(f, x, tp)
     xt = ReverseDiff.TrackedReal(x, zero(x), tp)
 
     y = f(x)
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
 
     yt = f(xt)
     @test yt == y
     dual = f(Dual(x, one(x)))
-    @test length(tp) == 1
-    instruction = first(tp)
+    recorded = take_recorded!(tp)
+    @test length(recorded) == 1
+    instruction = first(recorded)
     @test typeof(instruction) <: ReverseDiff.ScalarInstruction
     @test instruction.input === xt
     @test instruction.output === yt
     @test instruction.cache[] === partials(dual, 1)
-    empty!(tp)
     return nothing
 end
 
@@ -62,37 +62,37 @@ function test_forward(f, a, b, tp)
 
     c = f(a, b)
     dual = f(Dual(a, one(a), zero(a)), Dual(b, zero(b), one(b)))
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
 
     tc = f(at, b)
     @test tc == c
-    @test length(tp) == 1
-    instruction = first(tp)
+    recorded = take_recorded!(tp)
+    @test length(recorded) == 1
+    instruction = first(recorded)
     @test typeof(instruction) <: ReverseDiff.ScalarInstruction
     @test instruction.input === (at, b)
     @test instruction.output === tc
     @test instruction.cache[] === SVector(partials(dual, 1), partials(dual, 1))
-    empty!(tp)
 
     tc = f(a, bt)
     @test tc == c
-    @test length(tp) == 1
-    instruction = first(tp)
+    recorded = take_recorded!(tp)
+    @test length(recorded) == 1
+    instruction = first(recorded)
     @test typeof(instruction) <: ReverseDiff.ScalarInstruction
     @test instruction.input === (a, bt)
     @test instruction.output === tc
     @test instruction.cache[] === SVector(partials(dual, 2), partials(dual, 2))
-    empty!(tp)
 
     tc = f(at, bt)
     @test tc == c
-    @test length(tp) == 1
-    instruction = first(tp)
+    recorded = take_recorded!(tp)
+    @test length(recorded) == 1
+    instruction = first(recorded)
     @test typeof(instruction) <: ReverseDiff.ScalarInstruction
     @test instruction.input === (at, bt)
     @test instruction.output === tc
     @test instruction.cache[] === SVector(partials(dual)...)
-    empty!(tp)
     return nothing
 end
 
@@ -142,11 +142,11 @@ function test_skip(g, x, tp)
     xt = ReverseDiff.TrackedReal(x, zero(x), tp)
 
     y = g(x)
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
 
     yt = g(xt)
     @test yt === y
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
     return nothing
 end
 
@@ -155,19 +155,19 @@ function test_skip(g, a, b, tp)
     bt = ReverseDiff.TrackedReal(b, zero(b), tp)
 
     c = g(a, b)
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
 
     tc = g(at, b)
     @test tc === c
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
 
     tc = g(a, bt)
     @test tc === c
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
 
     tc = g(at, bt)
     @test tc === c
-    @test isempty(tp)
+    @test isempty(take_recorded!(tp))
     return nothing
 end