More thread-friendly compilation API method (#55)

allow compilation API to be used in thread-friendly way, and remove thread-unfriendly convenience methods

Author: jrevels

Diff for: benchmark/benchmark.jl

@@ -30,23 +30,26 @@ function grad_benchmark_driver!(out, f, x)
     gc()
 
     if length(tp.tape) <= 10000
-        ctp = ReverseDiff.compile(tp)
-
-        # warmup
-        ReverseDiff.gradient!(out, ctp, x)
-        ReverseDiff.forward_pass!(ctp)
-        ReverseDiff.reverse_pass!(ctp)
-
-        # actual
-        print("  gradient! (compiled): ")
-        @time ReverseDiff.gradient!(out, ctp, x)
-        gc()
-        print("  forward pass (compiled): ")
-        @time ReverseDiff.forward_pass!(ctp)
-        gc()
-        print("  reverse pass (compiled): ")
-        @time ReverseDiff.reverse_pass!(ctp)
-        gc()
+        @eval begin
+            out, x = $out, $x
+            ctp = ReverseDiff.compile($tp)
+
+            # warmup
+            ReverseDiff.gradient!(out, ctp, x)
+            ReverseDiff.forward_pass!(ctp)
+            ReverseDiff.reverse_pass!(ctp)
+
+            # actual
+            print("  gradient! (compiled): ")
+            @time ReverseDiff.gradient!(out, ctp, x)
+            gc()
+            print("  forward pass (compiled): ")
+            @time ReverseDiff.forward_pass!(ctp)
+            gc()
+            print("  reverse pass (compiled): ")
+            @time ReverseDiff.reverse_pass!(ctp)
+            gc()
+        end
     else
         println("skipping compiled GradientTape benchmark because the tape is too long ($(length(tp.tape)) elements)")
     end

Diff for: docs/src/api.md

@@ -9,23 +9,20 @@ CurrentModule = ReverseDiff
 ```@docs
 ReverseDiff.gradient
 ReverseDiff.gradient!
-ReverseDiff.compile_gradient
 ```
 
 ## Jacobians of `f(x::AbstractArray{<:Real}...)::AbstractArray{<:Real}`
 
 ```@docs
 ReverseDiff.jacobian
 ReverseDiff.jacobian!
-ReverseDiff.compile_jacobian
 ```
 
 ## Hessians of `f(x::AbstractArray{<:Real})::Real`
 
 ```@docs
 ReverseDiff.hessian
 ReverseDiff.hessian!
-ReverseDiff.compile_hessian
 ```
 
 ## The `AbstractTape` API

Diff for: examples/gradient.jl

@@ -1,4 +1,4 @@
-using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile_gradient
+using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile
 
 #########
 # setup #
@@ -7,8 +7,11 @@ using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile_gr
 # some objective function to work with
 f(a, b) = sum(a' * b + a * b')
 
-# generate a gradient function for `f` using inputs of shape 100x100 with Float64 elements
-const ∇f! = compile_gradient(f, (rand(100, 100), rand(100, 100)))
+# pre-record a GradientTape for `f` using inputs of shape 100x100 with Float64 elements
+const f_tape = GradientTape(f, (rand(100, 100), rand(100, 100)))
+
+# compile `f_tape` into a more optimized representation
+const compiled_f_tape = compile(f_tape)
 
 # some inputs and work buffers to play around with
 a, b = rand(100, 100), rand(100, 100)
@@ -20,16 +23,19 @@ cfg = GradientConfig(inputs)
 # taking gradients #
 ####################
 
-# with a pre-recorded/comiled tape (generated in the setup above) #
+# with pre-recorded/compiled tapes (generated in the setup above) #
 #-----------------------------------------------------------------#
+
 # this should be the fastest method, and non-allocating
+gradient!(results, compiled_f_tape, inputs)
 
-∇f!(results, inputs)
+# this should be the second fastest method, and also non-allocating
+gradient!(results, f_tape, inputs)
 
 # with a pre-allocated GradientConfig #
 #-------------------------------------#
-# this is more flexible than a pre-recorded tape, but can be wasteful since the tape
-# will be re-recorded for every call.
+# these methods are more flexible than a pre-recorded tape, but can be
+# wasteful since the tape will be re-recorded for every call.
 
 gradient!(results, f, inputs, cfg)
 

Diff for: examples/jacobian.jl

@@ -8,9 +8,13 @@ using ReverseDiff: JacobianTape, JacobianConfig, jacobian, jacobian!, compile_ja
 f(a, b) = (a + b) * (a * b)'
 g!(out, a, b) = A_mul_Bc!(out, a + b, a * b)
 
-# generate a jacobian function for `f` using inputs of shape 10x10 with Float64 elements
-const Jf! = compile_jacobian(f, (rand(10, 10), rand(10, 10)))
-const Jg! = compile_jacobian(g!, rand(10, 10), (rand(10, 10), rand(10, 10)))
+# pre-record JacobianTapes for `f` and `g` using inputs of shape 10x10 with Float64 elements
+const f_tape = JacobianTape(f, (rand(10, 10), rand(10, 10)))
+const g_tape = JacobianTape(g!, rand(10, 10), (rand(10, 10), rand(10, 10)))
+
+# compile `f_tape` and `g_tape` into more optimized representations
+const compiled_f_tape = compile(f_tape)
+const compiled_g_tape = compile(g_tape)
 
 # some inputs and work buffers to play around with
 a, b = rand(10, 10), rand(10, 10)
@@ -24,13 +28,16 @@ gcfg = JacobianConfig(output, inputs)
 # taking Jacobians #
 ####################
 
-# with a pre-recorded/comiled tape (generated in the setup above) #
+# with pre-recorded/compiled tapes (generated in the setup above) #
 #-----------------------------------------------------------------#
-# this should be the fastest method, and non-allocating
 
-Jf!(results, inputs)
+# these should be the fastest methods, and non-allocating
+jacobian!(results, compiled_f_tape, inputs)
+jacobian!(results, compiled_g_tape, inputs)
 
-Jg!(results, inputs)
+# these should be the second fastest methods, and also non-allocating
+jacobian!(results, f_tape, inputs)
+jacobian!(results, g_tape, inputs)
 
 # with a pre-allocated JacobianConfig #
 #-------------------------------------#

Diff for: src/api/gradients.jl

@@ -20,7 +20,7 @@ call.
 """
 function gradient(f, input, cfg::GradientConfig = GradientConfig(input))
     tape = GradientTape(f, input, cfg)
-    result = construct_result(tape.input)
+    result = construct_result(input_hook(tape))
     seeded_reverse_pass!(result, tape)
     empty!(cfg.tape)
     return result
@@ -59,7 +59,7 @@ If `input` is a tuple of `AbstractArray`s, assume `tape` represents a function o
 of `f` w.r.t. `input[i].`
 """
 function gradient!(tape::Union{GradientTape,CompiledGradient}, input)
-    result = construct_result(tape.input)
+    result = construct_result(input_hook(tape))
     gradient!(result, tape, input)
     return result
 end

Diff for: src/api/hessians.jl

@@ -67,7 +67,7 @@ Assuming `tape` represents a function of the form `f(::AbstractArray{<:Real})::R
 return the Hessian `H(f)(input)`.
 """
 function hessian!(tape::Union{HessianTape,CompiledHessian}, input::AbstractArray)
-    result = construct_result(tape.output, tape.input)
+    result = construct_result(output_hook(tape), input_hook(tape))
     hessian!(result, tape, input)
     return result
 end
@@ -92,7 +92,7 @@ end
 function hessian!(result::DiffResult, tape::Union{HessianTape,CompiledHessian}, input::AbstractArray)
     seeded_forward_pass!(tape, input)
     seeded_reverse_pass!(DiffResult(DiffBase.gradient(result), DiffBase.hessian(result)), tape)
-    DiffBase.value!(result, tape.func(input))
+    DiffBase.value!(result, func_hook(tape)(input))
     return result
 end
 

Diff for: src/api/jacobians.jl

@@ -60,7 +60,7 @@ function jacobian(f!, output, input, cfg::JacobianConfig = JacobianConfig(output
     tape = JacobianTape(f!, output, input, cfg)
     isa(input, TrackedArray) && empty!(input.tape)
     result = jacobian!(tape, input)
-    extract_result_value!(output, tape.output)
+    extract_result_value!(output, output_hook(tape))
     empty!(tape.tape)
     return result
 end
@@ -75,7 +75,7 @@ function jacobian!(result, f!, output, input, cfg::JacobianConfig = JacobianConf
     tape = JacobianTape(f!, output, input, cfg)
     isa(input, TrackedArray) && empty!(input.tape)
     jacobian!(result, tape, input)
-    extract_result_value!(output, tape.output)
+    extract_result_value!(output, output_hook(tape))
     empty!(tape.tape)
     return result
 end
@@ -102,7 +102,7 @@ execute `tape` with new `input` values. There is no way to re-run `tape`'s tape
 new `output` values into the tape.
 """
 function jacobian!(tape::Union{JacobianTape,CompiledJacobian}, input)
-    result = construct_result(tape.output, tape.input)
+    result = construct_result(output_hook(tape), input_hook(tape))
     jacobian!(result, tape, input)
     return result
 end
@@ -138,32 +138,32 @@ differentation.
 
 # function jacobian(f, input, cfg::JacobianConfig = JacobianConfig(input))
 #     tape = JacobianTape(f, input, cfg)
-#     result = construct_result(tape.output, tape.input)
-#     seeded_reverse_pass!(result, tape.output, tape.input, tape.tape)
+#     result = construct_result(output_hook(tape), input_hook(tape))
+#     seeded_reverse_pass!(result, output_hook(tape), input_hook(tape), tape.tape)
 #     empty!(tape.tape)
 #     return result
 # end
 #
 # function jacobian!(result, f, input, cfg::JacobianConfig = JacobianConfig(input))
 #     tape = JacobianTape(f, input, cfg)
-#     seeded_reverse_pass!(result, tape.output, tape.input, tape.tape)
+#     seeded_reverse_pass!(result, output_hook(tape), input_hook(tape), tape.tape)
 #     empty!(tape.tape)
 #     return result
 # end
 #
 # function jacobian(f!, output, input, cfg::JacobianConfig = JacobianConfig(output, input))
 #     tape = JacobianTape(f!, output, input, cfg)
-#     result = construct_result(tape.output, tape.input)
-#     seeded_reverse_pass!(result, tape.output, tape.input, tape.tape)
-#     extract_result_value!(output, tape.output)
+#     result = construct_result(output_hook(tape), input_hook(tape))
+#     seeded_reverse_pass!(result, output_hook(tape), input_hook(tape), tape.tape)
+#     extract_result_value!(output, output_hook(tape))
 #     empty!(tape.tape)
 #     return result
 # end
 #
 # function jacobian!(result, f!, output, input, cfg::JacobianConfig = JacobianConfig(output, input))
 #     tape = JacobianTape(f!, output, input, cfg)
-#     seeded_reverse_pass!(result, tape.output, tape.input, tape.tape)
-#     extract_result_value!(output, tape.output)
+#     seeded_reverse_pass!(result, output_hook(tape), input_hook(tape), tape.tape)
+#     extract_result_value!(output, output_hook(tape))
 #     empty!(tape.tape)
 #     return result
 # end